#!/usr/bin/env python3
"""
Add wettest/driest precipitation month metrics to data/climate-data.csv.

The metrics use the existing monthly NOAA nClimGrid precipitation file. They
are categorical month labels derived from 1991-2020 county-area means.
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

import numpy as np
import xarray as xr

from build_county_climate_data import (
    MONTH_NAMES,
    _as_monthly_climatology,
    _extract_grid_2d,
    _load_counties,
    _select_data_var,
    _zonal_mean,
)


REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CLIMATE_DATA = REPO_ROOT / "data" / "climate-data.csv"
DEFAULT_COUNTIES_GEOJSON = REPO_ROOT / "data" / "geojson-counties-fips.json"
DEFAULT_MONTHLY_PRCP_NC = REPO_ROOT / "data" / "noaa" / "nclimgrid" / "nclimgrid_prcp.nc"
PRECIP_MONTH_FIELDS = ["wettestPrecipMonth", "driestPrecipMonth"]


def normalize_fips(value: object) -> str:
    """Return a 5-digit county FIPS string from mixed text or numeric input."""
    digits = "".join(ch for ch in str(value).strip() if ch.isdigit())
    return digits.zfill(5)[-5:] if digits else ""


def read_csv_rows(path: Path) -> tuple[list[str], list[dict]]:
    """Read a CSV while preserving the source field order."""
    with path.open("r", encoding="utf-8-sig", newline="") as csv_file:
        reader = csv.DictReader(csv_file)
        if reader.fieldnames is None:
            raise ValueError(f"{path} has no CSV header.")
        return list(reader.fieldnames), list(reader)


def fieldnames_with_precip_months(original_fields: list[str]) -> list[str]:
    """Place the new fields next to the existing precipitation metrics."""
    fields = [field for field in original_fields if field not in PRECIP_MONTH_FIELDS]
    insert_after = "seasonalityIndex"
    insert_at = fields.index(insert_after) + 1 if insert_after in fields else len(fields)
    return fields[:insert_at] + PRECIP_MONTH_FIELDS + fields[insert_at:]


def build_precip_month_lookup(
    *,
    counties_geojson: Path,
    monthly_prcp_nc: Path,
    climatology_start_year: int,
    climatology_end_year: int,
) -> dict[str, tuple[str, str]]:
    """Calculate wettest and driest precipitation month for each county."""
    counties = _load_counties(counties_geojson)
    monthly_prcp = xr.open_dataset(monthly_prcp_nc, decode_times=True)
    try:
        prcp_var = _select_data_var(monthly_prcp, "mlyprcp_norm")
        prcp_monthly = _as_monthly_climatology(
            monthly_prcp[prcp_var],
            start_year=climatology_start_year,
            end_year=climatology_end_year,
        )

        prcp_by_month: list[list[float]] = []
        for month in range(12):
            prcp_arr, prcp_transform = _extract_grid_2d(prcp_monthly.isel(time=month))
            prcp_by_month.append(_zonal_mean(prcp_arr, prcp_transform, counties))
    finally:
        monthly_prcp.close()

    lookup: dict[str, tuple[str, str]] = {}
    for idx, row in counties.iterrows():
        county_fips = str(row["county_fips"])
        prcp_months_mm = np.array([prcp_by_month[m][idx] for m in range(12)], dtype=np.float64)
        if not np.any(np.isfinite(prcp_months_mm)):
            continue

        wettest_month = MONTH_NAMES[int(np.nanargmax(prcp_months_mm))]
        driest_month = MONTH_NAMES[int(np.nanargmin(prcp_months_mm))]
        lookup[county_fips] = (wettest_month, driest_month)

    return lookup


def apply_precipitation_month_metrics(
    *,
    climate_data: Path,
    counties_geojson: Path,
    monthly_prcp_nc: Path,
    out: Path,
    climatology_start_year: int,
    climatology_end_year: int,
) -> None:
    """Merge precipitation month categories into the app climate CSV."""
    original_fields, rows = read_csv_rows(climate_data)
    lookup = build_precip_month_lookup(
        counties_geojson=counties_geojson,
        monthly_prcp_nc=monthly_prcp_nc,
        climatology_start_year=climatology_start_year,
        climatology_end_year=climatology_end_year,
    )

    updated_count = 0
    missing_count = 0
    for row in rows:
        county_fips = normalize_fips(row.get("countyFips", ""))
        values = lookup.get(county_fips)
        if values:
            row["wettestPrecipMonth"], row["driestPrecipMonth"] = values
            updated_count += 1
        else:
            row["wettestPrecipMonth"] = ""
            row["driestPrecipMonth"] = ""
            missing_count += 1

        source = (row.get("source") or "").strip() or "unknown-source"
        tag = f"precip-month-extremes=noaa-nclimgrid-monthly-{climatology_start_year}-{climatology_end_year}"
        if tag not in source:
            row["source"] = f"{source} + {tag}"

    fieldnames = fieldnames_with_precip_months(original_fields)
    with out.open("w", encoding="utf-8", newline="") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)

    print(f"Wrote {len(rows)} rows to {out}")
    print(f"Updated precipitation month metrics for {updated_count} rows.")
    print(f"Left precipitation month metrics blank for {missing_count} rows.")


def parse_args() -> argparse.Namespace:
    """Parse command-line paths for the precipitation month update."""
    parser = argparse.ArgumentParser(
        description="Add wettestPrecipMonth and driestPrecipMonth to climate-data.csv."
    )
    parser.add_argument("--climate-data", type=Path, default=DEFAULT_CLIMATE_DATA)
    parser.add_argument("--counties-geojson", type=Path, default=DEFAULT_COUNTIES_GEOJSON)
    parser.add_argument("--monthly-prcp-nc", type=Path, default=DEFAULT_MONTHLY_PRCP_NC)
    parser.add_argument("--out", type=Path, default=DEFAULT_CLIMATE_DATA)
    parser.add_argument("--climatology-start-year", type=int, default=1991)
    parser.add_argument("--climatology-end-year", type=int, default=2020)
    return parser.parse_args()


def main() -> None:
    """Run the CSV update."""
    args = parse_args()
    apply_precipitation_month_metrics(
        climate_data=args.climate_data,
        counties_geojson=args.counties_geojson,
        monthly_prcp_nc=args.monthly_prcp_nc,
        out=args.out,
        climatology_start_year=args.climatology_start_year,
        climatology_end_year=args.climatology_end_year,
    )


if __name__ == "__main__":
    main()