#!/usr/bin/env python3
"""
Aggregate downloaded gridMET humidity/heat files to county-level metrics.

Expected input layout:
data/gridmet/<year>/sph.nc
data/gridmet/<year>/rmax.nc
data/gridmet/<year>/rmin.nc

Default output:
data/gridmet/county_gridmet_humidity.csv
"""

from __future__ import annotations

import argparse
import csv
import math
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path

import geopandas as gpd
import numpy as np
import pandas as pd
import xarray as xr


STATE_FIPS_TO_ABBR = {
    "01": "AL",
    "04": "AZ",
    "05": "AR",
    "06": "CA",
    "08": "CO",
    "09": "CT",
    "10": "DE",
    "11": "DC",
    "12": "FL",
    "13": "GA",
    "16": "ID",
    "17": "IL",
    "18": "IN",
    "19": "IA",
    "20": "KS",
    "21": "KY",
    "22": "LA",
    "23": "ME",
    "24": "MD",
    "25": "MA",
    "26": "MI",
    "27": "MN",
    "28": "MS",
    "29": "MO",
    "30": "MT",
    "31": "NE",
    "32": "NV",
    "33": "NH",
    "34": "NJ",
    "35": "NM",
    "36": "NY",
    "37": "NC",
    "38": "ND",
    "39": "OH",
    "40": "OK",
    "41": "OR",
    "42": "PA",
    "44": "RI",
    "45": "SC",
    "46": "SD",
    "47": "TN",
    "48": "TX",
    "49": "UT",
    "50": "VT",
    "51": "VA",
    "53": "WA",
    "54": "WV",
    "55": "WI",
    "56": "WY",
}
CONUS_STATE_FIPS = set(STATE_FIPS_TO_ABBR)
REQUIRED_VARIABLES = ("sph", "rmax", "rmin")
SUMMER_MONTHS = {6, 7, 8}
NOAA_REGION_CODE_TO_FIPS_OVERRIDES = {
    "18511": "11001",
}
NOAA_TMAX_SOURCE_FIPS_OVERRIDES = {
    "46113": (
        "46102",
        "Shannon County, SD changed name/code to Oglala Lakota County, SD effective 2015-05-01.",
    ),
    "51515": (
        "51019",
        "Bedford independent city, VA changed to town status and was added to Bedford County effective 2013-07-01.",
    ),
    "51678": (
        "51163",
        "Lexington independent city, VA uses surrounding Rockbridge County, VA as NOAA tmax proxy because NOAA county daily files do not include Lexington city.",
    ),
}


@dataclass(frozen=True)
class CountyGridMap:
    counties: gpd.GeoDataFrame
    pixel_indices: np.ndarray
    county_indices: np.ndarray
    weights: np.ndarray
    sorted_pixel_indices: np.ndarray
    sorted_county_indices: np.ndarray
    sorted_weights: np.ndarray
    group_starts: np.ndarray
    group_county_indices: np.ndarray
    lat_dim: str
    lon_dim: str


@dataclass
class YearData:
    datasets: list[xr.Dataset]
    arrays: dict[str, xr.DataArray]


def _normalize_fips(value: object, width: int = 5) -> str:
    text = str(value).strip()
    digits = "".join(ch for ch in text if ch.isdigit())
    return digits.zfill(width)[-width:] if digits else ""


def _load_state_crosswalk(path: Path) -> dict[str, str]:
    if not path.exists():
        print(
            f"State crosswalk not found at {path}. Falling back to the first two "
            "NOAA region-code digits as state FIPS."
        )
        return {}

    rows = list(csv.reader(path.open("r", encoding="utf-8-sig", newline="")))
    if not rows:
        return {}

    header = [cell.strip().lower() for cell in rows[0]]
    ncei_idx = next((idx for idx, cell in enumerate(header) if "ncei" in cell and "code" in cell), None)
    fips_idx = next((idx for idx, cell in enumerate(header) if "fips" in cell and "code" in cell), None)
    data_rows = rows[1:] if ncei_idx is not None and fips_idx is not None else rows

    mapping: dict[str, str] = {}
    for row in data_rows:
        if ncei_idx is not None and fips_idx is not None:
            if len(row) <= max(ncei_idx, fips_idx):
                continue
            ncei = _normalize_fips(row[ncei_idx], 2)
            fips = _normalize_fips(row[fips_idx], 2)
        else:
            codes = [_normalize_fips(cell, 2) for cell in row]
            codes = [code for code in codes if code]
            if len(codes) < 2:
                continue
            ncei, fips = codes[0], codes[1]
        if ncei and fips:
            mapping[ncei] = fips
    return mapping


def _county_fips_from_noaa_region(region_code: object, state_crosswalk: dict[str, str]) -> str:
    code = _normalize_fips(region_code, 5)
    if code in NOAA_REGION_CODE_TO_FIPS_OVERRIDES:
        return NOAA_REGION_CODE_TO_FIPS_OVERRIDES[code]
    state_fips = state_crosswalk.get(code[:2], code[:2])
    return f"{state_fips}{code[-3:]}" if state_fips else ""


def _load_counties(path: Path) -> gpd.GeoDataFrame:
    counties = gpd.read_file(path)
    if counties.crs is None:
        counties = counties.set_crs("EPSG:4326")
    else:
        counties = counties.to_crs("EPSG:4326")

    if "id" in counties.columns:
        fips = counties["id"]
    elif "GEOID" in counties.columns:
        fips = counties["GEOID"]
    elif "GEOID10" in counties.columns:
        fips = counties["GEOID10"]
    else:
        raise ValueError("County GeoJSON needs an id, GEOID, or GEOID10 column.")

    counties["countyFips"] = fips.map(_normalize_fips)
    counties["stateFips"] = counties["countyFips"].str.slice(0, 2)
    counties = counties[counties["stateFips"].isin(CONUS_STATE_FIPS)].copy()
    counties["countyName"] = counties.get("NAME", counties["countyFips"]).astype(str)
    counties["state"] = counties["stateFips"].map(STATE_FIPS_TO_ABBR)
    counties = counties.sort_values("countyFips").reset_index(drop=True)
    counties["countyIndex"] = np.arange(len(counties), dtype=np.int32)
    return counties


def _find_dimension(dataset: xr.Dataset, candidates: tuple[str, ...]) -> str:
    for candidate in candidates:
        if candidate in dataset.dims or candidate in dataset.coords:
            return candidate
    lower_lookup = {name.lower(): name for name in set(dataset.dims) | set(dataset.coords)}
    for candidate in candidates:
        if candidate.lower() in lower_lookup:
            return lower_lookup[candidate.lower()]
    raise ValueError(f"Unable to find one of these dimensions: {', '.join(candidates)}")


def _select_variable(dataset: xr.Dataset, preferred: str) -> str:
    if preferred in dataset.data_vars:
        return preferred
    for name in dataset.data_vars:
        if name.lower() == preferred.lower():
            return name
    if len(dataset.data_vars) == 1:
        return next(iter(dataset.data_vars))
    raise ValueError(f"Unable to select variable {preferred}; found {list(dataset.data_vars)}")


def _year_folder(path: Path) -> int | None:
    try:
        return int(path.name)
    except ValueError:
        return None


def _discover_complete_years(gridmet_dir: Path) -> list[int]:
    years: list[int] = []
    for child in sorted(gridmet_dir.iterdir()):
        if not child.is_dir():
            continue
        year = _year_folder(child)
        if year is None:
            continue
        if all((child / f"{variable}.nc").exists() for variable in REQUIRED_VARIABLES):
            years.append(year)
    return years


def _load_noaa_tmax_month(
    noaa_tmax_dir: Path,
    year: int,
    month: int,
    state_crosswalk: dict[str, str],
) -> dict[str, list[float]]:
    path = noaa_tmax_dir / str(year) / f"tmax-{year}{month:02d}-cty-scaled.csv"
    if not path.exists():
        return {}

    values: dict[str, list[float]] = {}
    with path.open("r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.reader(handle)
        for row in reader:
            if len(row) < 7 or row[0].strip().lower() != "cty":
                continue
            county_fips = _county_fips_from_noaa_region(row[1], state_crosswalk)
            day_values: list[float] = []
            for raw_value in row[6:]:
                try:
                    value = float(raw_value)
                except ValueError:
                    day_values.append(math.nan)
                    continue
                day_values.append(value if value > -999 else math.nan)
            values[county_fips] = day_values
    return values


def _noaa_tmax_for_date(
    *,
    date: pd.Timestamp,
    county_fips: Sequence[str],
    noaa_tmax_dir: Path,
    state_crosswalk: dict[str, str],
    month_cache: dict[tuple[int, int], dict[str, list[float]]],
) -> np.ndarray:
    cache_key = (int(date.year), int(date.month))
    if cache_key not in month_cache:
        month_cache[cache_key] = _load_noaa_tmax_month(
            noaa_tmax_dir,
            int(date.year),
            int(date.month),
            state_crosswalk,
        )

    month_values = month_cache[cache_key]
    result = np.full(len(county_fips), np.nan, dtype=np.float64)
    day_offset = int(date.day) - 1
    for index, fips in enumerate(county_fips):
        source_fips = NOAA_TMAX_SOURCE_FIPS_OVERRIDES.get(fips, (fips, ""))[0]
        values = month_values.get(source_fips)
        if values is not None and day_offset < len(values):
            result[index] = values[day_offset]
    return result


def _build_county_grid_map(counties: gpd.GeoDataFrame, sample_file: Path) -> CountyGridMap:
    with xr.open_dataset(sample_file) as dataset:
        lat_dim = _find_dimension(dataset, ("lat", "latitude", "y"))
        lon_dim = _find_dimension(dataset, ("lon", "longitude", "x"))
        lat_values = np.asarray(dataset[lat_dim].values)
        lon_values = np.asarray(dataset[lon_dim].values)

    if lat_values.ndim != 1 or lon_values.ndim != 1:
        raise ValueError("This script expects 1D latitude and longitude coordinates.")

    lon_grid, lat_grid = np.meshgrid(lon_values, lat_values)
    flat_lats = lat_grid.ravel()
    flat_lons = lon_grid.ravel()
    flat_indices = np.arange(flat_lats.size, dtype=np.int64)

    bounds = counties.total_bounds
    in_bounds = (
        (flat_lons >= bounds[0] - 0.25)
        & (flat_lats >= bounds[1] - 0.25)
        & (flat_lons <= bounds[2] + 0.25)
        & (flat_lats <= bounds[3] + 0.25)
    )
    point_frame = gpd.GeoDataFrame(
        {
            "pixelIndex": flat_indices[in_bounds],
            "lat": flat_lats[in_bounds],
        },
        geometry=gpd.points_from_xy(flat_lons[in_bounds], flat_lats[in_bounds]),
        crs="EPSG:4326",
    )

    joined = gpd.sjoin(
        point_frame,
        counties[["countyIndex", "geometry"]],
        how="inner",
        predicate="within",
    )
    joined = joined.drop_duplicates(subset=["pixelIndex", "countyIndex"])

    pixel_indices = joined["pixelIndex"].to_numpy(dtype=np.int64)
    county_indices = joined["countyIndex"].to_numpy(dtype=np.int32)
    weights = np.cos(np.deg2rad(joined["lat"].to_numpy(dtype=np.float64)))

    mapped_counties = set(county_indices.tolist())
    fallback_pixels: list[int] = []
    fallback_counties: list[int] = []
    fallback_weights: list[float] = []
    for county_index, geometry in enumerate(counties.geometry):
        if county_index in mapped_counties or geometry.is_empty:
            continue
        point = geometry.representative_point()
        lon_position = int(np.argmin(np.abs(lon_values - point.x)))
        lat_position = int(np.argmin(np.abs(lat_values - point.y)))
        flat_index = lat_position * len(lon_values) + lon_position
        fallback_pixels.append(flat_index)
        fallback_counties.append(county_index)
        fallback_weights.append(float(math.cos(math.radians(float(lat_values[lat_position])))))

    if fallback_pixels:
        pixel_indices = np.concatenate([pixel_indices, np.asarray(fallback_pixels, dtype=np.int64)])
        county_indices = np.concatenate([county_indices, np.asarray(fallback_counties, dtype=np.int32)])
        weights = np.concatenate([weights, np.asarray(fallback_weights, dtype=np.float64)])

    print(
        f"mapped {len(np.unique(county_indices)):,} counties to {len(pixel_indices):,} grid cells "
        f"({len(fallback_pixels):,} nearest-cell fallback counties)"
    )
    sort_order = np.argsort(county_indices, kind="stable")
    sorted_county_indices = county_indices[sort_order]
    group_starts = np.r_[0, np.flatnonzero(np.diff(sorted_county_indices)) + 1]
    group_county_indices = sorted_county_indices[group_starts]
    return CountyGridMap(
        counties=counties,
        pixel_indices=pixel_indices,
        county_indices=county_indices,
        weights=weights,
        sorted_pixel_indices=pixel_indices[sort_order],
        sorted_county_indices=sorted_county_indices,
        sorted_weights=weights[sort_order],
        group_starts=group_starts,
        group_county_indices=group_county_indices,
        lat_dim=lat_dim,
        lon_dim=lon_dim,
    )


def _county_means(
    data_array: xr.DataArray,
    time_index: int,
    time_dim: str,
    county_map: CountyGridMap,
) -> np.ndarray:
    data = data_array.isel({time_dim: time_index}).transpose(county_map.lat_dim, county_map.lon_dim)
    flat = np.asarray(data.values, dtype=np.float64).ravel()
    values = flat[county_map.pixel_indices]
    valid = np.isfinite(values)
    county_count = len(county_map.counties)
    sums = np.bincount(
        county_map.county_indices[valid],
        weights=values[valid] * county_map.weights[valid],
        minlength=county_count,
    )
    weight_sums = np.bincount(
        county_map.county_indices[valid],
        weights=county_map.weights[valid],
        minlength=county_count,
    )
    means = np.full(county_count, np.nan, dtype=np.float64)
    np.divide(sums, weight_sums, out=means, where=weight_sums > 0)
    return means


def _county_means_chunk(
    data_array: xr.DataArray,
    start_index: int,
    end_index: int,
    time_dim: str,
    county_map: CountyGridMap,
) -> np.ndarray:
    data = data_array.isel({time_dim: slice(start_index, end_index)}).transpose(
        time_dim,
        county_map.lat_dim,
        county_map.lon_dim,
    )
    raw = np.asarray(data.values, dtype=np.float64)
    selected = raw.reshape(raw.shape[0], -1)[:, county_map.sorted_pixel_indices]
    valid = np.isfinite(selected)

    weighted_values = np.where(valid, selected * county_map.sorted_weights, 0.0)
    weight_values = np.where(valid, county_map.sorted_weights, 0.0)
    sums = np.add.reduceat(weighted_values, county_map.group_starts, axis=1)
    weight_sums = np.add.reduceat(weight_values, county_map.group_starts, axis=1)

    grouped_means = np.full_like(sums, np.nan, dtype=np.float64)
    np.divide(sums, weight_sums, out=grouped_means, where=weight_sums > 0)

    means = np.full((raw.shape[0], len(county_map.counties)), np.nan, dtype=np.float64)
    means[:, county_map.group_county_indices] = grouped_means
    return means


def _heat_index_f(t_f: np.ndarray, rh_pct: np.ndarray) -> np.ndarray:
    rh = np.clip(rh_pct, 0, 100)
    heat_index = (
        -42.379
        + 2.04901523 * t_f
        + 10.14333127 * rh
        - 0.22475541 * t_f * rh
        - 0.00683783 * t_f * t_f
        - 0.05481717 * rh * rh
        + 0.00122874 * t_f * t_f * rh
        + 0.00085282 * t_f * rh * rh
        - 0.00000199 * t_f * t_f * rh * rh
    )

    low_rh_adjustment = ((13 - rh) / 4) * np.sqrt(np.maximum((17 - np.abs(t_f - 95)) / 17, 0))
    high_rh_adjustment = ((rh - 85) / 10) * ((87 - t_f) / 5)
    heat_index = np.where((rh < 13) & (80 <= t_f) & (t_f <= 112), heat_index - low_rh_adjustment, heat_index)
    heat_index = np.where((rh > 85) & (80 <= t_f) & (t_f <= 87), heat_index + high_rh_adjustment, heat_index)
    return np.where(t_f >= 80, heat_index, t_f)


def _time_dimension(data_array: xr.DataArray, lat_dim: str, lon_dim: str) -> str:
    candidates = [dim for dim in data_array.dims if dim not in {lat_dim, lon_dim}]
    if not candidates:
        raise ValueError(f"No time dimension found for {data_array.name}.")
    return candidates[0]


def _open_year_datasets(gridmet_dir: Path, year: int) -> YearData:
    datasets: list[xr.Dataset] = []
    arrays: dict[str, xr.DataArray] = {}
    for variable in REQUIRED_VARIABLES:
        dataset = xr.open_dataset(gridmet_dir / str(year) / f"{variable}.nc")
        datasets.append(dataset)
        variable_name = _select_variable(dataset, variable)
        arrays[variable] = dataset[variable_name]
    return YearData(datasets=datasets, arrays=arrays)


def _close_year_datasets(year_data: YearData) -> None:
    for dataset in year_data.datasets:
        try:
            dataset.close()
        except Exception:
            pass


def _round_or_blank(value: float, digits: int) -> str:
    if not np.isfinite(value):
        return ""
    return str(round(float(value), digits))


def summarize(args: argparse.Namespace) -> None:
    years = args.years or _discover_complete_years(args.gridmet_dir)
    if not years:
        raise FileNotFoundError(
            f"No complete gridMET years found in {args.gridmet_dir}. "
            "Run scripts/download_gridmet_data.py first."
        )

    if args.years is None:
        start = args.start_year if args.start_year is not None else 1991
        end = args.end_year if args.end_year is not None else 2020
        years = [year for year in years if start <= year <= end]
    elif args.start_year is not None or args.end_year is not None:
        start = args.start_year if args.start_year is not None else min(years)
        end = args.end_year if args.end_year is not None else max(years)
        years = [year for year in years if start <= year <= end]

    if not years:
        raise FileNotFoundError("No complete gridMET years matched the requested year filters.")

    sample_file = args.gridmet_dir / str(years[0]) / "sph.nc"
    counties = _load_counties(args.counties_geojson)
    county_map = _build_county_grid_map(counties, sample_file)
    county_count = len(county_map.counties)
    county_fips = county_map.counties["countyFips"].astype(str).tolist()
    state_crosswalk = _load_state_crosswalk(args.state_crosswalk)
    noaa_month_cache: dict[tuple[int, int], dict[str, list[float]]] = {}

    annual_sph_sum = np.zeros(county_count, dtype=np.float64)
    annual_sph_count = np.zeros(county_count, dtype=np.float64)
    summer_sph_sum = np.zeros(county_count, dtype=np.float64)
    summer_sph_count = np.zeros(county_count, dtype=np.float64)
    annual_rh_sum = np.zeros(county_count, dtype=np.float64)
    annual_rh_count = np.zeros(county_count, dtype=np.float64)
    summer_rh_sum = np.zeros(county_count, dtype=np.float64)
    summer_rh_count = np.zeros(county_count, dtype=np.float64)
    heat_day_sum = np.zeros(county_count, dtype=np.float64)
    years_with_heat_data = np.zeros(county_count, dtype=np.float64)

    for year in years:
        print(f"process {year}")
        year_data = _open_year_datasets(args.gridmet_dir, year)
        arrays = year_data.arrays
        try:
            time_dim = _time_dimension(arrays["sph"], county_map.lat_dim, county_map.lon_dim)
            dates = pd.to_datetime(arrays["sph"][time_dim].values)
            yearly_heat_days = np.zeros(county_count, dtype=np.float64)
            yearly_heat_valid_days = np.zeros(county_count, dtype=np.float64)

            for chunk_start in range(0, len(dates), args.chunk_days):
                chunk_end = min(chunk_start + args.chunk_days, len(dates))
                chunk_dates = dates[chunk_start:chunk_end]
                sph_chunk = _county_means_chunk(arrays["sph"], chunk_start, chunk_end, time_dim, county_map) * 1000
                rmax_chunk = _county_means_chunk(arrays["rmax"], chunk_start, chunk_end, time_dim, county_map)
                rmin_chunk = _county_means_chunk(arrays["rmin"], chunk_start, chunk_end, time_dim, county_map)

                for offset, date in enumerate(chunk_dates):
                    sph = sph_chunk[offset]
                    rmax = rmax_chunk[offset]
                    rmin = rmin_chunk[offset]
                    tmax_f = _noaa_tmax_for_date(
                        date=date,
                        county_fips=county_fips,
                        noaa_tmax_dir=args.noaa_tmax_dir,
                        state_crosswalk=state_crosswalk,
                        month_cache=noaa_month_cache,
                    ) * 9 / 5 + 32

                    rh_mean = (rmax + rmin) / 2
                    annual_sph_valid = np.isfinite(sph)
                    annual_rh_valid = np.isfinite(rh_mean)
                    annual_sph_sum[annual_sph_valid] += sph[annual_sph_valid]
                    annual_sph_count[annual_sph_valid] += 1
                    annual_rh_sum[annual_rh_valid] += rh_mean[annual_rh_valid]
                    annual_rh_count[annual_rh_valid] += 1

                    if int(date.month) in SUMMER_MONTHS:
                        summer_sph_sum[annual_sph_valid] += sph[annual_sph_valid]
                        summer_sph_count[annual_sph_valid] += 1
                        summer_rh_sum[annual_rh_valid] += rh_mean[annual_rh_valid]
                        summer_rh_count[annual_rh_valid] += 1

                    heat_valid = np.isfinite(tmax_f) & np.isfinite(rmin)
                    heat_index = _heat_index_f(tmax_f, rmin)
                    yearly_heat_days[heat_valid] += heat_index[heat_valid] >= args.heat_index_threshold_f
                    yearly_heat_valid_days[heat_valid] += 1

            has_heat_year = yearly_heat_valid_days > 0
            heat_day_sum[has_heat_year] += yearly_heat_days[has_heat_year]
            years_with_heat_data[has_heat_year] += 1
        finally:
            _close_year_datasets(year_data)

    avg_sph = np.divide(annual_sph_sum, annual_sph_count, out=np.full(county_count, np.nan), where=annual_sph_count > 0)
    avg_summer_sph = np.divide(
        summer_sph_sum,
        summer_sph_count,
        out=np.full(county_count, np.nan),
        where=summer_sph_count > 0,
    )
    avg_rh = np.divide(annual_rh_sum, annual_rh_count, out=np.full(county_count, np.nan), where=annual_rh_count > 0)
    avg_summer_rh = np.divide(
        summer_rh_sum,
        summer_rh_count,
        out=np.full(county_count, np.nan),
        where=summer_rh_count > 0,
    )
    humid_heat_days = np.divide(
        heat_day_sum,
        years_with_heat_data,
        out=np.full(county_count, np.nan),
        where=years_with_heat_data > 0,
    )
    source_period = f"{min(years)}-{max(years)}" if len(years) > 1 else str(years[0])

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=[
                "countyFips",
                "countyName",
                "state",
                "avgSpecificHumidityGKg",
                "avgSummerSpecificHumidityGKg",
                "avgRelHumidityPct",
                "avgSummerRelHumidityPct",
                "humidHeatDays",
                "humidHeatSourceFips",
                "humidHeatFipsAdjustment",
                "gridmetAnalysisYears",
                "gridmetSource",
            ],
        )
        writer.writeheader()
        for index, county in county_map.counties.iterrows():
            writer.writerow(
                {
                    "countyFips": county["countyFips"],
                    "countyName": county["countyName"],
                    "state": county["state"],
                    "avgSpecificHumidityGKg": _round_or_blank(avg_sph[index], 3),
                    "avgSummerSpecificHumidityGKg": _round_or_blank(avg_summer_sph[index], 3),
                    "avgRelHumidityPct": _round_or_blank(avg_rh[index], 2),
                    "avgSummerRelHumidityPct": _round_or_blank(avg_summer_rh[index], 2),
                    "humidHeatDays": _round_or_blank(humid_heat_days[index], 2),
                    "humidHeatSourceFips": NOAA_TMAX_SOURCE_FIPS_OVERRIDES.get(
                        county["countyFips"],
                        (county["countyFips"], ""),
                    )[0],
                    "humidHeatFipsAdjustment": NOAA_TMAX_SOURCE_FIPS_OVERRIDES.get(
                        county["countyFips"],
                        ("", ""),
                    )[1],
                    "gridmetAnalysisYears": len(years),
                    "gridmetSource": (
                        f"gridmet-{source_period} county-cell-weighted; "
                        "relative-humidity=(rmax+rmin)/2; "
                        "humidHeatDays=noaa-nclimgrid-tmax+gridmet-rmin heat-index-gte-90f"
                    ),
                }
            )

    print(f"wrote {args.output}")


def _parse_years(text: str) -> list[int]:
    years: set[int] = set()
    for part in text.split(","):
        token = part.strip()
        if not token:
            continue
        if "-" in token:
            start_text, end_text = token.split("-", 1)
            years.update(range(int(start_text), int(end_text) + 1))
        else:
            years.add(int(token))
    return sorted(years)


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Summarize downloaded gridMET humidity files into county-level CSV metrics."
    )
    parser.add_argument("--gridmet-dir", type=Path, default=Path("data/gridmet"))
    parser.add_argument("--noaa-tmax-dir", type=Path, default=Path("data/noaa/tmax_cty_scaled"))
    parser.add_argument("--state-crosswalk", type=Path, default=Path("data/noaa/nclimgrid/us-state-codes_ncei-to-fips.csv"))
    parser.add_argument("--counties-geojson", type=Path, default=Path("data/geojson-counties-fips.json"))
    parser.add_argument("--output", type=Path, default=Path("data/gridmet/county_gridmet_humidity.csv"))
    parser.add_argument("--years", type=_parse_years, default=None, help="Optional comma/range years, e.g. 1991-2020.")
    parser.add_argument("--start-year", type=int, default=None)
    parser.add_argument("--end-year", type=int, default=None)
    parser.add_argument("--heat-index-threshold-f", type=float, default=90.0)
    parser.add_argument("--chunk-days", type=int, default=31)
    summarize(parser.parse_args())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())