import dataclasses
import json
import logging
import os
import tempfile
from pathlib import Path
from typing import Callable, Dict, List, Tuple, Optional, Union

import numpy as np
from osgeo import gdal
from te_schemas import reporting
from te_schemas.aoi import AOI
from te_schemas.datafile import DataFile
from te_schemas.error_recode import ErrorRecodePolygons
from te_schemas.jobs import Job
from te_schemas.results import (
    URI,
    Band,
    DataType,
    JsonResults,
    Raster,
    RasterFileType,
    RasterResults,
    TiledRaster,
    VectorFalsePositive,
    VectorResults,
    VectorType,
)

from .. import workers
from ..util import accumulate_dicts, save_vrt, wkt_geom_to_geojson_file_string
from ..util_numba import zonal_total
from . import config
from .land_deg_numba import (
    recode_block_stats,
    recode_indicator_errors,
    sdg_status_expanded,
)
from .models import DegradationErrorRecodeSummaryParams, SummaryTableLDErrorRecode
from .worker import DegradationSummary

# Index mapping for the 4x4 crosstab arrays produced by recode_block_stats.
# Index 0: degraded (-1), 1: stable (0), 2: improved (1), 3: nodata (-32768)
_CROSSTAB_IDX_TO_CLASS = [-1, 0, 1, int(config.NODATA_VALUE)]

logger = logging.getLogger(__name__)


def rasterize_error_recode(
    out_file: Path,
    model_file: Path,
    geojson: ErrorRecodePolygons,
) -> None:
    # Convert layer into an integer code so that all three recode_deg_to
    # options can be encoded within a single tiff

    recode_to_trans_code = geojson.recode_to_trans_code_dict
    error_recode_dict = ErrorRecodePolygons.Schema().dump(geojson)

    for feat in error_recode_dict["features"]:
        feat["properties"]["error_recode"] = recode_to_trans_code[
            (
                feat["properties"]["recode_deg_to"],
                feat["properties"]["recode_stable_to"],
                feat["properties"]["recode_imp_to"],
            )
        ]

        # Convert periods_affected to a bitmask for rasterization
        # bit 1 = baseline, bit 2 = report_1, bit 4 = report_2
        periods_affected = feat["properties"].get("periods_affected")
        if not periods_affected:
            raise ValueError(
                f"periods_affected is required and cannot be empty for feature with "
                f"recode_deg_to={feat['properties'].get('recode_deg_to')}, "
                f"recode_stable_to={feat['properties'].get('recode_stable_to')}, "
                f"recode_imp_to={feat['properties'].get('recode_imp_to')}"
            )

        periods_mask = 0
        if "baseline" in periods_affected:
            periods_mask |= 1
        if "report_1" in periods_affected:
            periods_mask |= 2
        if "report_2" in periods_affected:
            periods_mask |= 4
        feat["properties"]["periods_mask"] = periods_mask

    # TODO: Assumes WGS84 for now
    # Create temporary files for each band
    temp_error_recode = tempfile.NamedTemporaryFile(
        suffix="_temp_error_recode.tif", delete=False
    ).name
    temp_periods_mask = tempfile.NamedTemporaryFile(
        suffix="_temp_periods_mask.tif", delete=False
    ).name

    try:
        rasterize_worker_error = workers.Rasterize(
            Path(temp_error_recode),
            model_file,
            error_recode_dict,
            "error_recode",
        )
        rasterize_worker_error.work()

        rasterize_worker_periods = workers.Rasterize(
            Path(temp_periods_mask),
            model_file,
            error_recode_dict,
            "periods_mask",
        )
        rasterize_worker_periods.work()

        # Combine both single-band rasters into one multi-band VRT
        # Note: Do NOT delete the temp files here as the VRT references them
        gdal.BuildVRT(
            str(out_file), [temp_error_recode, temp_periods_mask], separate=True
        )

    except Exception:
        # Only clean up temporary files if there was an error
        if os.path.exists(temp_error_recode):
            os.remove(temp_error_recode)
        if os.path.exists(temp_periods_mask):
            os.remove(temp_periods_mask)
        raise


def _process_block(
    params: DegradationErrorRecodeSummaryParams,
    in_array,
    mask,
    xoff: int,
    yoff: int,
    cell_areas_raw,
) -> Tuple[SummaryTableLDErrorRecode, List[Dict]]:
    """
    Process a block of data to recode multiple SDG bands and compute zonal statistics.

    Returns summary statistics for baseline and reporting periods, plus arrays to write.
    """
    # Extract band arrays
    baseline_array = None
    recode_array = in_array[params.band_dict["recode_bandnum"] - 1, :, :]
    periods_mask_array = in_array[params.band_dict["periods_mask_bandnum"] - 1, :, :]

    # Get baseline band
    if params.baseline_band_num is not None:
        baseline_array = in_array[params.band_dict["baseline_bandnum"] - 1, :, :]

    # Collect reporting period input arrays
    reporting_arrays = []
    for i, _band_num in enumerate(params.report_band_nums or []):
        key = f"reporting_{i + 1}_bandnum"
        reporting_arrays.append(in_array[params.band_dict[key] - 1, :, :])

    cell_areas = np.repeat(cell_areas_raw, mask.shape[1], axis=1)

    # Get the recode parameters
    codes, deg_to, stable_to, imp_to = params.trans_code_lists

    write_arrays = []
    summaries = {}

    # Process each band that exists
    if baseline_array is not None:
        # recode_indicator_errors currently supports up to 2 reporting periods
        reporting_1_input = reporting_arrays[0] if len(reporting_arrays) > 0 else None
        reporting_2_input = reporting_arrays[1] if len(reporting_arrays) > 1 else None

        baseline_recoded, reporting_1_recoded, reporting_2_recoded = (
            recode_indicator_errors(
                baseline_array,
                reporting_1_input,
                reporting_2_input,
                recode_array,
                periods_mask_array,
                np.array(codes, dtype=np.int16),
                np.array(deg_to, dtype=np.int16),
                np.array(stable_to, dtype=np.int16),
                np.array(imp_to, dtype=np.int16),
            )
        )

        # Collect non-None recoded reporting arrays into a list
        recoded_reports = []
        for idx in range(len(reporting_arrays)):
            recoded = (
                reporting_1_recoded
                if idx == 0
                else reporting_2_recoded
                if idx == 1
                else None
            )
            if recoded is not None:
                recoded_reports.append(recoded)

        # Compute baseline summary and all crosstabs in a single fused pass.
        # Stack recoded reports into a 3D array (n_reports, rows, cols).
        if recoded_reports:
            reports_3d = np.stack(recoded_reports, axis=0)
        else:
            reports_3d = np.empty(
                (0, baseline_recoded.shape[0], baseline_recoded.shape[1]),
                dtype=np.int16,
            )

        baseline_totals, crosstabs_3d = recode_block_stats(
            baseline_recoded, reports_3d, cell_areas, mask
        )

        # Convert baseline totals array to dict for existing summary pipeline
        summaries["baseline_summary"] = {
            cls: float(baseline_totals[i])
            for i, cls in enumerate(_CROSSTAB_IDX_TO_CLASS)
            if baseline_totals[i] > 0
        }

        # Store crosstab 4x4 arrays as a list (one per reporting period)
        if len(recoded_reports) > 0:
            summaries["crosstabs"] = [
                crosstabs_3d[k] for k in range(crosstabs_3d.shape[0])
            ]

        # Add baseline to write arrays
        write_arrays.append(
            {
                "array": baseline_recoded,
                "xoff": xoff,
                "yoff": yoff,
                "band_name": "baseline",
            }
        )

        # Process each reporting period for status maps and summaries
        report_summaries = []
        for i, recoded in enumerate(recoded_reports):
            # Calculate 7-class status map
            status = sdg_status_expanded(baseline_recoded, recoded)
            report_summaries.append(zonal_total(status, cell_areas, mask))

            write_arrays.append(
                {
                    "array": status,
                    "xoff": xoff,
                    "yoff": yoff,
                    "band_name": f"reporting_{i + 1}_status",
                }
            )

            if params.write_reporting_sdg_tifs:
                write_arrays.append(
                    {
                        "array": recoded,
                        "xoff": xoff,
                        "yoff": yoff,
                        "band_name": f"reporting_{i + 1}_raw",
                    }
                )

        if report_summaries:
            summaries["report_summaries"] = report_summaries

    # Include the periods_mask array as the error recode output
    # This encodes which periods were affected:
    # 0 = no change, 1 = baseline, 2 = report_1, 3 = baseline+report_1,
    # 4 = report_2, 5 = baseline+report_2, 6 = report_1+report_2, 7 = all
    write_arrays.append(
        {"array": periods_mask_array, "xoff": xoff, "yoff": yoff, "band_name": "recode"}
    )

    # Create summary table
    summary_table = SummaryTableLDErrorRecode(
        baseline_summary=summaries.get("baseline_summary", {}),
        report_summaries=summaries.get("report_summaries"),
        crosstabs=summaries.get("crosstabs"),
    )

    return summary_table, write_arrays


def _accumulate_summary_tables(
    tables: List[SummaryTableLDErrorRecode],
) -> SummaryTableLDErrorRecode:
    if len(tables) == 1:
        return tables[0]
    else:
        out = tables[0]

        for table in tables[1:]:
            out.baseline_summary = accumulate_dicts(
                [out.baseline_summary, table.baseline_summary]
            )

            # Accumulate report summaries (list of dicts)
            if out.report_summaries is not None and table.report_summaries is not None:
                out.report_summaries = [
                    accumulate_dicts([a, b])
                    for a, b in zip(out.report_summaries, table.report_summaries)
                ]
            elif table.report_summaries is not None:
                out.report_summaries = table.report_summaries

            # Accumulate crosstab arrays with numpy addition
            if out.crosstabs is not None and table.crosstabs is not None:
                out.crosstabs = [a + b for a, b in zip(out.crosstabs, table.crosstabs)]
            elif table.crosstabs is not None:
                out.crosstabs = [ct.copy() for ct in table.crosstabs]

        return out


def _get_error_recode_input_vrt(baseline_df=None, reporting_dfs=None, error_df=None):
    """
    Create a VRT with multiple input bands for error recoding.

    Parameters:
    - baseline_df: DataFile for baseline SDG layer (optional)
    - reporting_dfs: List of DataFiles for reporting period SDG layers (optional)
    - error_df: DataFile containing error recode and periods_affected bands (required)

    Returns:
    - out_vrt: Path to the created VRT file
    - vrt_band_dict: Dictionary mapping band names to band numbers in the VRT
    """
    band_vrts = []
    df_band_list = []
    band_counter = 1

    # Add baseline if available
    if baseline_df is not None:
        band_vrts.append(save_vrt(baseline_df.path, 1))
        df_band_list.append(("baseline_bandnum", band_counter))
        band_counter += 1

    # Add reporting periods
    for i, rdf in enumerate(reporting_dfs or []):
        band_vrts.append(save_vrt(rdf.path, 1))
        df_band_list.append((f"reporting_{i + 1}_bandnum", band_counter))
        band_counter += 1

    # Add error recode band (required)
    if error_df is not None:
        band_vrts.append(
            save_vrt(
                error_df.path,
                error_df.index_for_name(config.ERROR_RECODE_INPUT_POLYS_BAND_NAME) + 1,
            )
        )
        df_band_list.append(("recode_bandnum", band_counter))
        band_counter += 1

        # Add periods_affected band if it exists
        try:
            periods_band_index = error_df.index_for_name("periods_affected")
            band_vrts.append(save_vrt(error_df.path, periods_band_index + 1))
            df_band_list.append(("periods_mask_bandnum", band_counter))
            band_counter += 1
        except (ValueError, AttributeError):
            # periods_affected band doesn't exist, create a dummy band
            # This ensures backward compatibility
            pass

    out_vrt = tempfile.NamedTemporaryFile(
        suffix="_error_recode_inputs.vrt", delete=False
    ).name
    gdal.BuildVRT(out_vrt, [vrt for vrt in band_vrts], separate=True)

    vrt_band_dict = {item[0]: item[1] for item in df_band_list}

    return out_vrt, vrt_band_dict


def _prepare_df(path, band_str, band_index) -> DataFile:
    band = Band(**band_str)

    return DataFile(path=Path(save_vrt(path, band_index)), bands=[band])


def _prepare_error_recode_df(path, band_str) -> DataFile:
    """Special version of _prepare_df for error recode files with multiple bands."""
    # Create the primary band (Error recode) - always band 1
    primary_band = Band(**band_str)

    # Open the raster to check for additional bands
    ds = gdal.Open(str(path))
    bands = [primary_band]

    if ds and ds.RasterCount >= 2:
        # Assume the second band is the periods_affected band
        # (since we create it that way in rasterize_error_recode)
        periods_band = Band(
            name="periods_affected",
            metadata={},
            no_data_value=primary_band.no_data_value,
            activated=True,
        )
        bands.append(periods_band)

    if ds:
        ds = None  # Close the dataset

    return DataFile(path=Path(path), bands=bands)


def _format_crosstab(crosstab_array, from_period, to_period):
    """Convert a 4x4 numpy crosstab array to the output format.

    The array is produced by ``recode_block_stats`` and indexed by SDG class:
        Index 0: degraded (-1), 1: stable (0), 2: improved (1), 3: nodata.

    Masked pixels are never counted (they are skipped in the numba loop),
    so no mask-value filtering is needed here.

    Args:
        crosstab_array: numpy float64 array of shape (4, 4) with areas in
            sq km.
        from_period: Name of the from period (e.g., 'baseline').
        to_period: Name of the to period (e.g., 'report_1').

    Returns:
        Dict with from_period, to_period, total_area_km2, and transitions
        matrix in the standard output format.
    """
    class_names = ["degraded", "stable", "improved", "nodata"]

    total_area = float(crosstab_array.sum())

    transitions = {}
    for from_idx, from_name in enumerate(class_names):
        transitions[from_name] = {}
        for to_idx, to_name in enumerate(class_names):
            area = float(crosstab_array[from_idx, to_idx])
            transitions[from_name][to_name] = {
                "area_km2": round(area, 3),
                "area_pct": round(area / total_area * 100, 2)
                if total_area > 0
                else 0.0,
            }

    return {
        "from_period": from_period,
        "to_period": to_period,
        "total_area_km2": round(total_area, 3),
        "transitions": transitions,
    }


def get_serialized_results(st, layer_name, periods_to_output=None):
    """Get serialized results for specified periods.

    Args:
        st: Summary table containing baseline_summary, report_summaries, crosstabs
        layer_name: Name for the output layer
        periods_to_output: List of periods to include in output. If None, includes all.
                          Valid values: 'baseline', 'report_1', 'report_2', ...
    """
    # Default to all periods if not specified
    if periods_to_output is None:
        n_reports = len(st.report_summaries) if st.report_summaries else 0
        periods_to_output = ["baseline"] + [f"report_{i + 1}" for i in range(n_reports)]

    def create_area_list(summary_dict, title_suffix=""):
        return reporting.AreaList(
            f"SDG Indicator 15.3.1{title_suffix}",
            "sq km",
            [
                reporting.Area("Improved", summary_dict.get(1, 0.0)),
                reporting.Area("Stable", summary_dict.get(0, 0.0)),
                reporting.Area("Degraded", summary_dict.get(-1, 0.0)),
                reporting.Area(
                    "No data", summary_dict.get(int(config.NODATA_VALUE), 0)
                ),
            ],
        )

    def create_status_area_list(status_summary_dict, title_suffix=""):
        """Create area list from 7-class status summary."""
        return reporting.AreaList(
            f"SDG Indicator 15.3.1{title_suffix}",
            "sq km",
            [
                # Degraded classes (1, 2, 3)
                reporting.Area(
                    "Persistent degradation", status_summary_dict.get(1, 0.0)
                ),
                reporting.Area("Recent degradation", status_summary_dict.get(2, 0.0)),
                reporting.Area("Baseline degradation", status_summary_dict.get(3, 0.0)),
                # Stable class (4)
                reporting.Area("Stability", status_summary_dict.get(4, 0.0)),
                # Improved classes (5, 6, 7)
                reporting.Area("Baseline improvement", status_summary_dict.get(5, 0.0)),
                reporting.Area("Recent improvement", status_summary_dict.get(6, 0.0)),
                reporting.Area(
                    "Persistent improvement", status_summary_dict.get(7, 0.0)
                ),
                # No data
                reporting.Area(
                    "No data", status_summary_dict.get(int(config.NODATA_VALUE), 0)
                ),
            ],
        )

    reports = []
    baseline_summary = None

    # Add baseline summary only if baseline is in periods_to_output
    if "baseline" in periods_to_output:
        baseline_summary = create_area_list(st.baseline_summary, " - Baseline")
        reports.append(baseline_summary)

    # Add reporting period summaries dynamically
    report_periods = [p for p in periods_to_output if p.startswith("report_")]
    for period_name in report_periods:
        # Extract 0-based index from period name (e.g. "report_1" -> 0)
        try:
            period_idx = int(period_name.split("_")[1]) - 1
        except (IndexError, ValueError):
            continue
        if st.report_summaries is not None and period_idx < len(st.report_summaries):
            report_summary = create_status_area_list(
                st.report_summaries[period_idx],
                f" - Reporting Period {period_idx + 1} Status",
            )
            reports.append(report_summary)

    # Create the main report - use baseline if available, otherwise first available summary
    if baseline_summary is not None:
        sdg_report = reporting.SDG15Report(summary=baseline_summary)
        additional_reports = reports[1:] if len(reports) > 1 else []
    elif len(reports) > 0:
        # No baseline in output, use first available report as primary
        sdg_report = reporting.SDG15Report(summary=reports[0])
        additional_reports = reports[1:] if len(reports) > 1 else []
    else:
        # No reports at all - create empty summary
        empty_summary = create_area_list({}, " - No Data")
        sdg_report = reporting.SDG15Report(summary=empty_summary)
        additional_reports = []

    # Add additional summaries to the serialized output
    result_dict = dataclasses.asdict(sdg_report)
    if len(additional_reports) > 0:
        result_dict["additional_summaries"] = [
            dataclasses.asdict(report) for report in additional_reports
        ]

    # Add crosstab data if available
    crosstab_results = []
    for period_name in report_periods:
        try:
            period_idx = int(period_name.split("_")[1]) - 1
        except (IndexError, ValueError):
            continue
        if st.crosstabs is not None and period_idx < len(st.crosstabs):
            crosstab_results.append(
                _format_crosstab(st.crosstabs[period_idx], "baseline", period_name)
            )
    if crosstab_results:
        if "report" not in result_dict:
            result_dict["report"] = {}
        result_dict["report"]["crosstabs"] = crosstab_results

    return result_dict


def recode_errors(params) -> Job:
    logger.debug("Entering recode_errors")
    aoi = AOI(params["aoi"])
    job_output_path = Path(params["output_path"])
    baseline_df = _prepare_df(
        params["layer_baseline_band_path"],
        params["layer_baseline_band"],
        params["layer_baseline_band_index"],
    )

    # Collect reporting period datafiles dynamically
    # NOTE: This loop expects SEQUENTIAL keys (layer_reporting_1_*, layer_reporting_2_*,
    # etc.) and breaks at the first gap. Callers MUST provide all intermediate
    # reporting period keys up to the highest one needed. If only report_2 is
    # needed, report_1 keys must also be present.
    reporting_dfs = []
    i = 1
    MAX_REPORTING_PERIODS_LOOKAHEAD = 5
    while True:
        key_path = f"layer_reporting_{i}_band_path"
        key_band = f"layer_reporting_{i}_band"
        key_index = f"layer_reporting_{i}_band_index"
        if key_path not in params:
            # Check if there are any higher-numbered reporting periods that
            # would be silently skipped (indicates a caller bug)
            for j in range(i + 1, i + MAX_REPORTING_PERIODS_LOOKAHEAD):
                if f"layer_reporting_{j}_band_path" in params:
                    logger.warning(
                        f"Gap in reporting period keys: layer_reporting_{i}_band_path "
                        f"missing but layer_reporting_{j}_band_path exists. "
                        f"Reporting periods after the gap will be ignored."
                    )
                    break
            break
        reporting_dfs.append(
            _prepare_df(params[key_path], params[key_band], params[key_index])
        )
        i += 1

    error_recode_df = _prepare_error_recode_df(
        params["layer_error_recode_path"],
        params["layer_error_recode_band"],
    )
    logger.debug("Loading error polygons")
    error_polygons_data = params["error_polygons"]
    error_polygons = ErrorRecodePolygons(
        features=error_polygons_data.get("features", []),
        name=error_polygons_data.get("name"),
        crs=error_polygons_data.get("crs"),
        type=error_polygons_data.get("type", "FeatureCollection"),
    )

    # Get periods to include in output (defaults to all if not specified)
    periods_to_output = params.get(
        "periods_to_output",
        ["baseline"] + [f"report_{i + 1}" for i in range(len(reporting_dfs))],
    )

    # Get option to include polygon GeoJSON in output (defaults to False)
    include_polygon_geojson = params.get("include_polygon_geojson", False)

    logger.debug("Running _compute_error_recode")
    summary_table, error_recode_paths = _compute_error_recode(
        baseline_df=baseline_df,
        reporting_dfs=reporting_dfs,
        error_recode_df=error_recode_df,
        error_polygons=error_polygons,
        job_output_path=job_output_path.parent / f"{job_output_path.stem}.json",
        aoi=aoi,
    )

    if params["write_tifs"]:
        # Create bands dynamically based on what was processed and what should be in output
        # Use standard band names for styling compatibility, with period info in metadata
        out_bands = []

        # Extract baseline years from the baseline band metadata
        baseline_band_metadata = params.get("layer_baseline_band", {}).get(
            "metadata", {}
        )
        baseline_year_initial = baseline_band_metadata.get("year_initial")
        baseline_year_final = baseline_band_metadata.get("year_final")

        # Include baseline band only if baseline is in periods_to_output
        if "baseline" in periods_to_output:
            baseline_metadata = (
                dict(params["metadata"]) if params.get("metadata") else {}
            )
            baseline_metadata["period"] = "Baseline"
            out_bands.append(
                Band(
                    name=config.SDG_BAND_NAME,
                    no_data_value=int(config.NODATA_VALUE),
                    metadata=baseline_metadata,
                    add_to_map=True,
                    activated=True,
                )
            )

        # Add reporting period bands dynamically
        for i, rdf in enumerate(reporting_dfs):
            period_name = f"report_{i + 1}"
            if period_name not in periods_to_output:
                continue
            # Get reporting period year info from its band metadata
            report_band_metadata = params.get(f"layer_reporting_{i + 1}_band", {}).get(
                "metadata", {}
            )
            period_metadata = {
                "period": f"Period {i + 1}",
                "baseline_year_initial": baseline_year_initial,
                "baseline_year_final": baseline_year_final,
                "reporting_year_initial": report_band_metadata.get("year_initial"),
                "reporting_year_final": report_band_metadata.get("year_final"),
            }
            out_bands.append(
                Band(
                    name=config.SDG_STATUS_BAND_NAME,
                    no_data_value=int(config.NODATA_VALUE),
                    metadata=period_metadata,
                    add_to_map=True,
                    activated=True,
                )
            )

            # Add raw reporting period band if write_reporting_sdg_tifs is enabled
            if params.get("write_reporting_sdg_tifs", False):
                period_raw_metadata = (
                    dict(report_band_metadata) if report_band_metadata else {}
                )
                period_raw_metadata["period"] = f"Period {i + 1}"
                period_raw_metadata["type"] = "raw"
                out_bands.append(
                    Band(
                        name=config.SDG_BAND_NAME,
                        no_data_value=int(config.NODATA_VALUE),
                        metadata=period_raw_metadata,
                        add_to_map=False,
                        activated=False,
                    )
                )

        # Always include the error recode band (periods affected)
        out_bands.append(
            Band(
                name=config.ERROR_RECODE_PERIODS_BAND_NAME,
                no_data_value=int(config.NODATA_VALUE),
                metadata=params["metadata"],  # copy metadata from input job
                add_to_map=False,
                activated=False,
            )
        )

        if len(error_recode_paths) > 1:
            error_recode_vrt = (
                job_output_path.parent / f"{job_output_path.stem}_error_recode.vrt"
            )
            gdal.BuildVRT(str(error_recode_vrt), [str(p) for p in error_recode_paths])
            rasters: Dict[str, Union[Raster, TiledRaster]] = {
                DataType.INT16.value: TiledRaster(
                    tile_uris=[URI(uri=p) for p in error_recode_paths],
                    uri=URI(uri=error_recode_vrt),
                    bands=out_bands,
                    datatype=DataType.INT16,
                    filetype=RasterFileType.COG,
                )
            }
            main_uri = URI(uri=error_recode_vrt)
        else:
            rasters: Dict[str, Union[Raster, TiledRaster]] = {
                DataType.INT16.value: Raster(
                    uri=URI(uri=error_recode_paths[0]),
                    bands=out_bands,
                    datatype=DataType.INT16,
                    filetype=RasterFileType.COG,
                )
            }
            main_uri = URI(uri=error_recode_paths[0])

        raster_results = RasterResults(
            name=params["layer_baseline_band"]["name"],
            uri=main_uri,
            rasters=rasters,
            data=get_serialized_results(
                summary_table,
                params["layer_baseline_band"]["name"] + " recode",
                periods_to_output=periods_to_output,
            ),
        )

        # If include_polygon_geojson is True, save the polygons to a GeoJSON file
        # and return both raster and vector results
        if include_polygon_geojson:
            # Write error polygons to GeoJSON file
            geojson_path = (
                job_output_path.parent
                / f"{job_output_path.stem}_error_polygons.geojson"
            )
            error_polygons_dict = ErrorRecodePolygons.Schema().dump(error_polygons)
            with open(geojson_path, "w", encoding="utf-8") as f:
                json.dump(error_polygons_dict, f, indent=2)

            logger.debug(f"Saved error recode polygons to {geojson_path}")

            # Create VectorResults for the error polygons
            vector_results = VectorResults(
                name="Error recode periods",
                vector=VectorFalsePositive(
                    uri=URI(uri=geojson_path),
                    type=VectorType.ERROR_RECODE,
                ),
            )

            # Return list of both results
            results = [raster_results, vector_results]
        else:
            results = raster_results

    else:
        json_results = JsonResults(
            name=params["layer_baseline_band"]["name"],
            data=get_serialized_results(
                summary_table,
                params["layer_baseline_band"]["name"] + " recode",
                periods_to_output=periods_to_output,
            ),
        )

        # If include_polygon_geojson is True for JsonResults, embed in data
        if include_polygon_geojson:
            json_results.data["error_recode_polygons"] = (
                ErrorRecodePolygons.Schema().dump(error_polygons)
            )

        results = json_results

    return results


def _compute_error_recode(
    baseline_df=None,
    reporting_dfs=None,
    error_recode_df=None,
    error_polygons=None,
    job_output_path=None,
    aoi=None,
    write_reporting_sdg_tifs=False,
    mask_worker_function: Optional[Callable] = None,
    mask_worker_params: Optional[dict] = None,
    error_recode_worker_function: Optional[Callable] = None,
    error_recode_worker_params: Optional[dict] = None,
):
    reporting_dfs = reporting_dfs or []

    in_vrt, band_dict = _get_error_recode_input_vrt(
        baseline_df, reporting_dfs, error_recode_df
    )

    wkt_aois = aoi.meridian_split(as_extent=False, out_format="wkt")
    # Use the first available datafile for getting bounds
    model_df = (
        baseline_df or (reporting_dfs[0] if reporting_dfs else None) or error_recode_df
    )
    bbs = aoi.get_aligned_output_bounds(model_df.path)

    error_recode_name_pattern = {
        1: f"{job_output_path.stem}" + "_error_recode.tif",
        2: f"{job_output_path.stem}" + "_error_recode_{index}.tif",
    }[len(bbs)]
    mask_name_fragment = {
        1: "Generating mask for error recoding",
        2: "Generating mask for error recoding (part {index} of 2)",
    }[len(wkt_aois)]

    error_recode_tables = []
    error_recode_paths = []
    error_message = None

    for index, (wkt_aoi, this_bbs) in enumerate(zip(wkt_aois, bbs), start=1):
        cropped_in_vrt = tempfile.NamedTemporaryFile(
            suffix="_ld_error_recode_inputs.vrt", delete=False
        ).name
        gdal.BuildVRT(cropped_in_vrt, in_vrt, outputBounds=this_bbs)

        mask_tif = tempfile.NamedTemporaryFile(
            suffix="_ld_error_recode_mask.tif", delete=False
        ).name
        logger.info(f"Saving mask to {mask_tif}")
        logger.info(
            str(job_output_path.parent / mask_name_fragment.format(index=index))
        )
        geojson = wkt_geom_to_geojson_file_string(wkt_aoi)

        if mask_worker_function:
            mask_result = mask_worker_function(
                Path(mask_tif),
                geojson,
                Path(cropped_in_vrt),
                **(mask_worker_params or {}),
            )
        else:
            mask_worker = workers.Mask(
                Path(mask_tif),
                geojson,
                Path(cropped_in_vrt),
            )
            mask_result = mask_worker.work()

        if mask_result:
            out_path = job_output_path.parent / error_recode_name_pattern.format(
                index=index
            )
            error_recode_paths.append(out_path)

            logger.info(f"Calculating error recode and saving layer to: {out_path}")

            # Determine which band numbers are available
            baseline_band_num = band_dict.get("baseline_bandnum")
            report_band_nums = [
                band_dict[f"reporting_{i + 1}_bandnum"]
                for i in range(len(reporting_dfs))
                if f"reporting_{i + 1}_bandnum" in band_dict
            ]

            # Calculate the number of output bands accurately
            n_out_bands = 1  # Always include recode band (periods_mask)

            # Add baseline band if available
            if baseline_df is not None:
                n_out_bands += 1

            # Add reporting period bands
            for _rdf in reporting_dfs:
                n_out_bands += 1  # Status map
                if write_reporting_sdg_tifs:
                    n_out_bands += 1  # Raw recoded band

            error_recode_params = DegradationErrorRecodeSummaryParams(
                in_file=str(cropped_in_vrt),
                out_file=str(out_path),
                band_dict=band_dict,
                model_band_number=1,
                n_out_bands=n_out_bands,
                mask_file=mask_tif,
                trans_code_lists=error_polygons.trans_code_lists,
                write_reporting_sdg_tifs=write_reporting_sdg_tifs,
                baseline_band_num=baseline_band_num,
                report_band_nums=report_band_nums if report_band_nums else None,
            )

            if error_recode_worker_function:
                result = error_recode_worker_function(
                    error_recode_params,
                    _process_block,
                    **(error_recode_worker_params or {}),
                )
            else:
                summarizer = DegradationSummary(error_recode_params, _process_block)
                result = summarizer.work()

            if not result:
                if hasattr(result, "is_killed") and result.is_killed():
                    error_message = "Cancelled calculation of error recoding."
                else:
                    error_message = "Failed while performing error recode."
                    result = None
            else:
                error_recode_tables.append(_accumulate_summary_tables(result))

        else:
            error_message = "Error creating mask."

    if error_message:
        logger.error(error_message)
        raise RuntimeError(f"Error during error recode: {error_message}")

    error_recode_table = _accumulate_summary_tables(error_recode_tables)

    return error_recode_table, error_recode_paths
