"""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 os
from pathlib import Path
from typing import Literal

import cyclopts
import geopandas as gpd
from rich.progress import track

DATA_DIR = Path(os.environ.get("DATA_DIR", "../../data")) / "entropyc-rts"
LEVEL2_PATH = DATA_DIR / "darts" / "DARTS_NitzeEtAl_v1-2_features_2018-2023_level2.parquet"
LEVEL2_COV_PATH = DATA_DIR / "darts" / "DARTS_NitzeEtAl_v1-2_coverage_2018-2023_level2.parquet"


def extract_darts_rts(grid: Literal["hex", "healpix"], level: int):
    """Extract RTS labels from DARTS dataset.

    Args:
        grid (Literal["hex", "healpix"]): The grid type to use.
        level (int): The grid level to use.

    """
    darts_l2 = gpd.read_parquet(LEVEL2_PATH)
    darts_cov_l2 = gpd.read_parquet(LEVEL2_COV_PATH)

    gridname = f"permafrost_{grid}{level}"
    grid = gpd.read_parquet(DATA_DIR / f"grids/{gridname}_grid.parquet")

    grid_l2 = grid.overlay(darts_l2.to_crs(grid.crs), how="intersection")
    grid_cov_l2 = grid.overlay(darts_cov_l2.to_crs(grid.crs), how="intersection")

    years = list(grid_cov_l2["year"].unique())
    for year in track(years, total=len(years), description="Processing years..."):
        subset = grid_l2[grid_l2["year"] == year]
        subset_cov = grid_cov_l2[grid_cov_l2["year"] == year]

        counts = subset.groupby("cell_id").size()
        grid[f"darts_{year}_rts_count"] = grid.cell_id.map(counts)

        areas = subset.groupby("cell_id").apply(lambda x: x.geometry.area.sum(), include_groups=False)
        grid[f"darts_{year}_rts_area"] = grid.cell_id.map(areas)

        areas_cov = subset_cov.groupby("cell_id").apply(lambda x: x.geometry.area.sum(), include_groups=False)
        grid[f"darts_{year}_covered_area"] = grid.cell_id.map(areas_cov)
        grid[f"darts_{year}_coverage"] = grid[f"darts_{year}_covered_area"] / grid.geometry.area

        grid[f"darts_{year}_rts_density"] = grid[f"darts_{year}_rts_area"] / grid[f"darts_{year}_covered_area"]

    output_path = DATA_DIR / f"darts/{gridname}_darts.parquet"
    grid.to_parquet(output_path)
    print(f"Saved RTS labels to {output_path}")


def main():  # noqa: D103
    cyclopts.run(extract_darts_rts)


if __name__ == "__main__":
    main()