"""Labels of Retrogressive-Thaw-Slumps (RTS).

Assumes that the level 1 and level 2 DARTS features have been downloaded into $DATA_DIR / entropyc-rts / darts: https://arcticdata.io/catalog/view/doi:10.18739/A22B8VD7C

Author: Tobias Hölzer
Date: October 2025
"""

import cyclopts
import geopandas as gpd
import pandas as pd
import xarray as xr
import xdggs
from rich import pretty, traceback
from stopuhr import stopwatch

from entropice.spatial import grids
from entropice.utils.paths import (
    DARTS_MLLABELS_DIR,
    DARTS_V1_DIR,
    DARTS_V2_DIR,
    get_darts_file,
)
from entropice.utils.types import Grid

traceback.install()
pretty.install()

darts_v2_l2_file = DARTS_V2_DIR / "all_prediction_segments_ensemble5.parquet"
darts_v2_l2_cov_file = DARTS_V2_DIR / "all_prediction_extent_ensemble5.parquet"
darts_v1_l2_file = DARTS_V1_DIR / "DARTS_NitzeEtAl_v1-2_features_2018-2023_level2.parquet"
darts_v1_l2_cov_file = DARTS_V1_DIR / "DARTS_NitzeEtAl_v1-2_coverage_2018-2023_level2.parquet"
darts_v1_corrections = DARTS_V1_DIR / "negative_correction.geojson"
darts_ml_training_labels_repo = DARTS_MLLABELS_DIR / "ML_training_labels" / "retrogressive_thaw_slumps"

cli = cyclopts.App(name="darts-rts")


def _load_grid(grid: Grid, level: int) -> tuple[gpd.GeoDataFrame, xr.DataArray]:
    grid_gdf = grids.open(grid, level)
    if grid == "hex":
        grid_gdf["cell_id"] = grid_gdf["cell_id"].apply(lambda x: int(x, 16))
    cell_areas = xr.DataArray(
        grid_gdf["cell_area"].to_numpy(), dims=["cell_id"], coords={"cell_id": grid_gdf["cell_id"].to_numpy()}
    )
    # We only want to use the geometry and the cell_id
    grid_gdf = grid_gdf[["cell_id", "geometry"]]
    return grid_gdf, cell_areas


@stopwatch("Convert xdgss")
def _convert_xdggs(darts: xr.Dataset, grid: Grid, level: int) -> xr.Dataset:
    darts = darts.rename({"cell_id": "cell_ids"})
    darts["cell_ids"] = darts["cell_ids"].astype("uint64")
    gridinfo = {
        "grid_name": "h3" if grid == "hex" else grid,
        "level": level,
    }
    if grid == "healpix":
        gridinfo["indexing_scheme"] = "nested"
    darts.cell_ids.attrs = gridinfo
    darts = xdggs.decode(darts)
    return darts


@stopwatch("Process RTS grid")
def _process_rts_grid(
    rts: gpd.GeoDataFrame,
    cov: gpd.GeoDataFrame,
    cell_areas: xr.DataArray,
) -> xr.Dataset:
    assert "cell_id" in rts.columns, "RTS data must contain cell_id column."
    assert "cell_id" in cov.columns, "Coverage data must contain cell_id column."

    cov["area_km2"] = cov.geometry.area / 1e6
    covered_area: pd.Series = cov.pivot_table(
        index="cell_id",
        values="area_km2",
        aggfunc="sum",
        fill_value=0,
    )["area_km2"]
    covered_area: xr.DataArray = xr.DataArray.from_series(covered_area)
    coverage = covered_area / cell_areas.sel(cell_id=covered_area.cell_id)
    rts["area_km2"] = rts.geometry.area / 1e6
    rts = rts.pivot_table(
        index="cell_id",
        values=["geometry", "area_km2"],
        fill_value=0,
        aggfunc={"geometry": "count", "area_km2": "sum"},  # type: ignore[arg-type]
    ).rename(columns={"geometry": "count"})
    rts: xr.Dataset = xr.Dataset.from_dataframe(rts)
    darts = xr.merge(
        [rts, covered_area.rename("covered_area_km2"), coverage.rename("coverage")], join="outer", fill_value=0
    )
    darts["density"] = (darts["area_km2"] / darts["covered_area_km2"]).fillna(0)
    return darts


@stopwatch("Process RTS yearly grid")
def _process_rts_yearly_grid(
    rts: gpd.GeoDataFrame,
    cov: gpd.GeoDataFrame,
    cell_areas: xr.DataArray,
) -> xr.Dataset:
    assert "cell_id" in rts.columns, "RTS data must contain cell_id column."
    assert "cell_id" in cov.columns, "Coverage data must contain cell_id column."

    cov["area_km2"] = cov.geometry.area / 1e6
    covered_area: pd.Series = cov.pivot_table(  # type: ignore[var-annotated]  # noqa: PD010
        index="cell_id",
        columns="year",
        values="area_km2",
        aggfunc="sum",
        fill_value=0,
    ).unstack()
    covered_area: xr.DataArray = xr.DataArray.from_series(covered_area)
    coverage = covered_area / cell_areas.sel(cell_id=covered_area.cell_id)
    rts: pd.DataFrame = (
        rts.pivot_table(  # type: ignore[var-annotated]  # noqa: PD013
            index="cell_id",
            columns="year",
            values=["geometry", "area_km2"],
            fill_value=0,
            aggfunc={"geometry": "count", "area_km2": "sum"},  # type: ignore[arg-type]
        )
        .stack(future_stack=True)
        .rename(columns={"geometry": "count"})
    )
    rts: xr.Dataset = xr.Dataset.from_dataframe(rts)
    darts = xr.merge(
        [rts, covered_area.rename("covered_area_km2"), coverage.rename("coverage")], join="outer", fill_value=0
    )
    darts["density"] = (darts["area_km2"] / darts["covered_area_km2"]).fillna(0)
    return darts


@cli.command()
def extract_darts_v2(grid: Grid, level: int):
    """Extract RTS labels from DARTS-v2 Level-2 dataset.

    Creates a Darts-v1 xarray Dataset on the specified grid and level.
    The Dataset contains the following variables:
    - count: Number of RTS in the cell
    - area_km2: Total area of RTS in the cell (in km^2)
    - covered_area_km2: Area of the cell covered by DARTS (in km^2)
    - coverage: Fraction of the cell covered by DARTS
    - density: Density of RTS area per covered area (area_km2 / covered_area_km2)
    Since the DARTS-v1 Level-2 dataset contains yearly data, all variables are indexed by year as well.
    Thus each variable has dimensions (cell_ids, year).

    Args:
        grid (Grid): The grid type to use.
        level (int): The grid level to use.

    """
    with stopwatch("Load data"):
        grid_gdf, cell_areas = _load_grid(grid, level)
        darts_l2 = gpd.read_parquet(darts_v2_l2_file).to_crs(grid_gdf.crs)
        darts_cov_l2 = gpd.read_parquet(darts_v2_l2_cov_file).to_crs(grid_gdf.crs)
        # Need to filter small noise pixels (I do not know where they are coming from)
        darts_cov_l2 = darts_cov_l2[darts_cov_l2.geometry.area > 1e9]

    with stopwatch("Assign RTS to grid"):
        grid_l2 = grid_gdf.overlay(darts_l2, how="intersection")
        grid_cov_l2 = grid_gdf.overlay(darts_cov_l2, how="intersection")

    darts = _process_rts_yearly_grid(grid_l2, grid_cov_l2, cell_areas)
    darts = _convert_xdggs(darts, grid, level)
    output_path = get_darts_file(grid, level, version="v2")
    with stopwatch(f"Writing Darts v2 to {output_path}"):
        darts.to_zarr(output_path, consolidated=False, mode="w")


@cli.command()
def extract_darts_v2_aggregated(grid: Grid, level: int):
    """Extract RTS labels from DARTS-v2 Level-3 dataset.

    Creates a Darts-v2 xarray Dataset on the specified grid and level.
    The Dataset contains the following variables:
    - count: Number of RTS in the cell
    - area_km2: Total area of RTS in the cell (in km^2)
    - covered_area_km2: Area of the cell covered by DARTS (in km^2)
    - coverage: Fraction of the cell covered by DARTS
    - density: Density of RTS area per covered area (area_km2 / covered_area_km2)
    Since the DARTS-v2 Level-2 dataset contains yearly data, the data is dissolved then exploded to obtain Level-3 data.
    Thus each variable has only the dimension (cell_ids).

    Args:
        grid (Grid): The grid type to use.
        level (int): The grid level to use.

    """
    with stopwatch("Load data"):
        grid_gdf, cell_areas = _load_grid(grid, level)
        darts_l2 = gpd.read_parquet(darts_v2_l2_file).to_crs(grid_gdf.crs)
        darts_cov_l2 = gpd.read_parquet(darts_v2_l2_cov_file).to_crs(grid_gdf.crs)
        # Need to filter small noise pixels (I do not know where they are coming from)
        darts_cov_l2 = darts_cov_l2[darts_cov_l2.geometry.area > 1e9]
        # Remove overlapping labels by dissolving
        darts_l2 = darts_l2[["geometry"]].dissolve().explode()
        darts_cov_l2 = darts_cov_l2[["geometry"]].dissolve().explode()

    with stopwatch("Extract RTS labels"):
        grid_l3 = grid_gdf.overlay(darts_l2, how="intersection")
        grid_cov_l3 = grid_gdf.overlay(darts_cov_l2, how="intersection")

    darts = _process_rts_grid(grid_l3, grid_cov_l3, cell_areas)
    darts = _convert_xdggs(darts, grid, level)
    output_path = get_darts_file(grid, level, version="v2-l3")
    with stopwatch(f"Writing Darts v2 l3 to {output_path}"):
        darts.to_zarr(output_path, consolidated=False, mode="w")


@cli.command()
def extract_darts_v1(grid: Grid, level: int):
    """Extract RTS labels from DARTS-v1 Level-2 dataset.

    Creates a Darts-v1 xarray Dataset on the specified grid and level.
    The Dataset contains the following variables:
    - count: Number of RTS in the cell
    - area_km2: Total area of RTS in the cell (in km^2)
    - covered_area_km2: Area of the cell covered by DARTS (in km^2)
    - coverage: Fraction of the cell covered by DARTS
    - density: Density of RTS area per covered area (area_km2 / covered_area_km2)
    Since the DARTS-v1 Level-2 dataset contains yearly data, all variables are indexed by year as well.
    Thus each variable has dimensions (cell_ids, year).

    Args:
        grid (Grid): The grid type to use.
        level (int): The grid level to use.

    """
    with stopwatch("Load data"):
        darts_l2 = gpd.read_parquet(darts_v1_l2_file)
        darts_cov_l2 = gpd.read_parquet(darts_v1_l2_cov_file)
        grid_gdf, cell_areas = _load_grid(grid, level)
        corrections = gpd.read_file(darts_v1_corrections).to_crs(darts_l2.crs)

    with stopwatch("Apply corrections"):
        # The correction file is just an area of sure negatives
        # Thus, we first need to remove all RTS labels that intersect with the correction area,
        darts_l2 = gpd.overlay(darts_l2, corrections, how="difference")
        # then we need to add the correction area as coverage to the coverage file per year.
        darts_cov_l2_cor = []
        for year in darts_cov_l2["year"].unique():
            year_cov = darts_cov_l2[darts_cov_l2["year"] == year]
            year_cov_cor = gpd.overlay(year_cov, corrections, how="union")
            year_cov_cor["year"] = year
            darts_cov_l2_cor.append(year_cov_cor)
        darts_cov_l2_cor = gpd.GeoDataFrame(pd.concat(darts_cov_l2_cor, ignore_index=True))
        darts_cov_l2_cor["year"] = darts_cov_l2_cor["year"].astype(int)
        darts_cov_l2 = darts_cov_l2_cor

    with stopwatch("Assign RTS to grid"):
        grid_l2 = grid_gdf.overlay(darts_l2.to_crs(grid_gdf.crs), how="intersection")
        grid_cov_l2 = grid_gdf.overlay(darts_cov_l2.to_crs(grid_gdf.crs), how="intersection")

    darts = _process_rts_yearly_grid(grid_l2, grid_cov_l2, cell_areas)
    darts = _convert_xdggs(darts, grid, level)
    output_path = get_darts_file(grid, level, version="v1")
    with stopwatch(f"Writing Darts v1 to {output_path}"):
        darts.to_zarr(output_path, consolidated=False, mode="w")


@cli.command()
def extract_darts_v1_aggregated(grid: Grid, level: int):
    """Extract RTS labels from DARTS-v1 Level-3 dataset.

    Creates a Darts-v1 xarray Dataset on the specified grid and level.
    The Dataset contains the following variables:
    - count: Number of RTS in the cell
    - area_km2: Total area of RTS in the cell (in km^2)
    - covered_area_km2: Area of the cell covered by DARTS (in km^2)
    - coverage: Fraction of the cell covered by DARTS
    - density: Density of RTS area per covered area (area_km2 / covered_area_km2)
    Since the DARTS-v1 Level-2 dataset contains yearly data, the data is dissolved then exploded to obtain Level-3 data.
    Thus each variable has only the dimension (cell_ids).

    Args:
        grid (Grid): The grid type to use.
        level (int): The grid level to use.

    """
    with stopwatch("Load data"):
        darts_l2 = gpd.read_parquet(darts_v1_l2_file)
        darts_cov_l2 = gpd.read_parquet(darts_v1_l2_cov_file)
        grid_gdf, cell_areas = _load_grid(grid, level)
        corrections = gpd.read_file(darts_v1_corrections).to_crs(darts_l2.crs)
        # Remove overlapping labels by dissolving
        darts_l2 = darts_l2[["geometry"]].dissolve().explode()
        darts_cov_l2 = darts_cov_l2[["geometry"]].dissolve().explode()

    with stopwatch("Apply corrections"):
        # The correction file is just an area of sure negatives
        # Thus, we first need to remove all RTS labels that intersect with the correction area,
        darts_l2 = gpd.overlay(darts_l2, corrections, how="difference")
        # then we need to add the correction area as coverage to the coverage file.
        darts_cov_l2 = gpd.overlay(darts_cov_l2, corrections, how="union")

    with stopwatch("Extract RTS labels"):
        grid_l3 = grid_gdf.overlay(darts_l2.to_crs(grid_gdf.crs), how="intersection")
        grid_cov_l3 = grid_gdf.overlay(darts_cov_l2.to_crs(grid_gdf.crs), how="intersection")

    darts = _process_rts_grid(grid_l3, grid_cov_l3, cell_areas)
    darts = _convert_xdggs(darts, grid, level)
    output_path = get_darts_file(grid, level, version="v1-l3")
    with stopwatch(f"Writing Darts v1 l3 to {output_path}"):
        darts.to_zarr(output_path, consolidated=False, mode="w")


@cli.command()
def extract_darts_mllabels(grid: Grid, level: int):
    """Extract RTS labels from the DARTS-mllabels dataset.

    Creates a Darts-mllabels xarray Dataset on the specified grid and level.
    The Dataset contains the following variables:
    - count: Number of RTS in the cell
    - area_km2: Total area of RTS in the cell (in km^2)
    - covered_area_km2: Area of the cell covered by DARTS (in km^2)
    - coverage: Fraction of the cell covered by DARTS
    - density: Density of RTS area per covered area (area_km2 / covered_area_km2)
    Since the DARTS-mllabels dataset contains not much data, all variables are aggregated over the entire time period.
    Thus each variable has only the dimension (cell_ids).

    Args:
        grid (Grid): The grid type to use.
        level (int): The grid level to use.

    """
    with stopwatch("Load data"):
        grid_gdf, cell_areas = _load_grid(grid, level)
        darts_mllabels = (
            gpd.GeoDataFrame(
                pd.concat([gpd.read_file(f) for f in darts_ml_training_labels_repo.glob("**/TrainingLabel*.gpkg")])
            )
            .reset_index(drop=True)
            .to_crs(grid_gdf.crs)
        )
        # Filter out invalid labels
        darts_mllabels = darts_mllabels[darts_mllabels.geometry.is_valid]
        # Remove overlapping labels by dissolving
        darts_mllabels = darts_mllabels[["geometry"]].dissolve().explode()

        darts_cov_mllabels = (
            gpd.GeoDataFrame(
                pd.concat([gpd.read_file(f) for f in darts_ml_training_labels_repo.glob("**/ImageFootprints*.gpkg")])
            )
            .reset_index(drop=True)
            .to_crs(grid_gdf.crs)
        )
        darts_cov_mllabels = darts_cov_mllabels[["geometry"]].dissolve().explode()

    with stopwatch("Extract RTS labels"):
        grid_mllabels = grid_gdf.overlay(darts_mllabels.to_crs(grid_gdf.crs), how="intersection")
        grid_cov_mllabels = grid_gdf.overlay(darts_cov_mllabels.to_crs(grid_gdf.crs), how="intersection")

    darts = _process_rts_grid(grid_mllabels, grid_cov_mllabels, cell_areas)
    darts = _convert_xdggs(darts, grid, level)
    output_path = get_darts_file(grid, level, version="mllabels")
    with stopwatch(f"Writing Darts v1 to {output_path}"):
        darts.to_zarr(output_path, consolidated=False, mode="w")


def main():  # noqa: D103
    cli()


if __name__ == "__main__":
    main()
