#!/usr/bin/env python3
"""
Build county-level average diurnal temperature range from NOAA nClimGrid-Daily.

The output is an ETL artifact used to keep climate-data.csv lean:

data/noaa/county_diurnal_temperature_range.csv

Metric definition:
- avgDiurnalTempRangeF: mean daily county Tmax - Tmin, expressed as a
  Fahrenheit temperature difference, over the selected analysis period.
"""

from __future__ import annotations

import argparse
import csv
from dataclasses import dataclass
from pathlib import Path

from build_county_locally_extreme_data import (
    DEFAULT_STATE_CROSSWALK,
    DEFAULT_TMAX_DIR,
    DEFAULT_TMIN_DIR,
    load_state_crosswalk,
    require_monthly_file,
    rows_by_fips,
)


REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUT = REPO_ROOT / "data" / "noaa" / "county_diurnal_temperature_range.csv"
SOURCE_TAG = "noaa-nclimgrid-daily-county-area-averages-scaled"


@dataclass
class DiurnalRangeTotals:
    """Running total for one county's daily Tmax-Tmin differences."""

    noaa_region_code: str = ""
    county_name: str = ""
    total_range_c: float = 0.0
    valid_days: int = 0
    skipped_negative_days: int = 0


def c_delta_to_f_delta(value_c: float) -> float:
    """Convert a Celsius temperature difference to a Fahrenheit difference."""
    return value_c * 9.0 / 5.0


def build_diurnal_temperature_range(
    *,
    start_year: int,
    end_year: int,
    tmax_dir: Path,
    tmin_dir: Path,
    state_crosswalk_path: Path,
) -> dict[str, DiurnalRangeTotals]:
    """Average daily county Tmax-Tmin over the requested analysis years."""
    state_crosswalk = load_state_crosswalk(state_crosswalk_path)
    totals: dict[str, DiurnalRangeTotals] = {}

    for year in range(start_year, end_year + 1):
        print(f"Reading daily Tmax/Tmin for {year}...")
        for month in range(1, 13):
            tmax_path = require_monthly_file(tmax_dir, "tmax", year, month)
            tmin_path = require_monthly_file(tmin_dir, "tmin", year, month)
            tmax_rows = rows_by_fips(
                tmax_path,
                year=year,
                month=month,
                state_crosswalk=state_crosswalk,
            )
            tmin_rows = rows_by_fips(
                tmin_path,
                year=year,
                month=month,
                state_crosswalk=state_crosswalk,
            )

            for county_fips in sorted(set(tmax_rows) & set(tmin_rows)):
                tmax_record = tmax_rows[county_fips]
                tmin_record = tmin_rows[county_fips]
                county_totals = totals.setdefault(
                    county_fips,
                    DiurnalRangeTotals(
                        noaa_region_code=tmax_record.region_code,
                        county_name=tmax_record.county_name,
                    ),
                )

                for tmax_value, tmin_value in zip(tmax_record.values, tmin_record.values):
                    if tmax_value is None or tmin_value is None:
                        continue
                    daily_range_c = tmax_value - tmin_value
                    if daily_range_c < 0:
                        county_totals.skipped_negative_days += 1
                        continue
                    county_totals.total_range_c += daily_range_c
                    county_totals.valid_days += 1

    return totals


def write_diurnal_temperature_range_csv(
    *,
    path: Path,
    totals: dict[str, DiurnalRangeTotals],
    start_year: int,
    end_year: int,
) -> None:
    """Write county DTR summary rows."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "countyFips",
        "noaaRegionCode",
        "countyName",
        "avgDiurnalTempRangeC",
        "avgDiurnalTempRangeF",
        "validDays",
        "skippedNegativeDays",
        "analysisStartYear",
        "analysisEndYear",
        "source",
    ]

    with path.open("w", encoding="utf-8", newline="") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        writer.writeheader()
        for county_fips in sorted(totals):
            county_totals = totals[county_fips]
            avg_range_c = (
                county_totals.total_range_c / county_totals.valid_days
                if county_totals.valid_days
                else None
            )
            writer.writerow(
                {
                    "countyFips": county_fips,
                    "noaaRegionCode": county_totals.noaa_region_code,
                    "countyName": county_totals.county_name,
                    "avgDiurnalTempRangeC": f"{avg_range_c:.2f}" if avg_range_c is not None else "",
                    "avgDiurnalTempRangeF": (
                        f"{c_delta_to_f_delta(avg_range_c):.2f}" if avg_range_c is not None else ""
                    ),
                    "validDays": county_totals.valid_days,
                    "skippedNegativeDays": county_totals.skipped_negative_days,
                    "analysisStartYear": start_year,
                    "analysisEndYear": end_year,
                    "source": SOURCE_TAG,
                }
            )

    print(f"Wrote {len(totals)} county DTR rows to {path}")


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Build county average diurnal temperature range from cached NOAA Tmax/Tmin files."
    )
    parser.add_argument("--start-year", type=int, default=1991)
    parser.add_argument("--end-year", type=int, default=2020)
    parser.add_argument("--tmax-dir", type=Path, default=DEFAULT_TMAX_DIR)
    parser.add_argument("--tmin-dir", type=Path, default=DEFAULT_TMIN_DIR)
    parser.add_argument("--state-crosswalk", type=Path, default=DEFAULT_STATE_CROSSWALK)
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    return parser.parse_args()


def main() -> None:
    """Run the DTR ETL."""
    args = parse_args()
    totals = build_diurnal_temperature_range(
        start_year=args.start_year,
        end_year=args.end_year,
        tmax_dir=args.tmax_dir,
        tmin_dir=args.tmin_dir,
        state_crosswalk_path=args.state_crosswalk,
    )
    write_diurnal_temperature_range_csv(
        path=args.out,
        totals=totals,
        start_year=args.start_year,
        end_year=args.end_year,
    )


if __name__ == "__main__":
    main()