#!/usr/bin/env python3
"""
Build county-level climate records for the web app filters.

Outputs a CSV file compatible with the browser app:
climate-data.csv

Metrics produced per county:
- koppenZone: majority Koppen-Geiger class
- avgTempF: annual mean temperature from NOAA 1991-2020 gridded normals
- annualPrecipIn: annual total precipitation from NOAA 1991-2020 gridded normals
- seasonalityIndex: precipitation seasonality, coefficient of variation of monthly totals (%)
- wettestPrecipMonth: month with the highest 1991-2020 county mean precipitation
- driestPrecipMonth: month with the lowest 1991-2020 county mean precipitation
- extremeDays: count of normal-days with Tmax >= hot threshold or Tmin <= freeze threshold
- avgSolarGhiKwhM2Day: annual average daily global horizontal irradiance (GHI), when a solar raster or representative-point CSV is provided

This script is intended for offline generation of complete county records.
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path
from typing import Dict, List, Tuple

import geopandas as gpd
import numpy as np
import rasterio
from rasterio.features import geometry_mask
import xarray as xr

DEFAULT_COUNTIES_GEOJSON_URL = "https://raw.githubusercontent.com/plotly/datasets/master/geojson-counties-fips.json"


STATE_FIPS_TO_ABBR = {
    "01": "AL",
    "02": "AK",
    "04": "AZ",
    "05": "AR",
    "06": "CA",
    "08": "CO",
    "09": "CT",
    "10": "DE",
    "11": "DC",
    "12": "FL",
    "13": "GA",
    "15": "HI",
    "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",
    "60": "AS",
    "66": "GU",
    "69": "MP",
    "72": "PR",
    "78": "VI",
}

# Beck et al legend key is expected as text file, but this default handles common codes.
DEFAULT_KOPPEN_CODE_MAP = {
    1: "Af",
    2: "Am",
    3: "Aw",
    4: "BWh",
    5: "BWk",
    6: "BSh",
    7: "BSk",
    8: "Csa",
    9: "Csb",
    10: "Csc",
    11: "Cwa",
    12: "Cwb",
    13: "Cwc",
    14: "Cfa",
    15: "Cfb",
    16: "Cfc",
    17: "Dsa",
    18: "Dsb",
    19: "Dsc",
    20: "Dsd",
    21: "Dwa",
    22: "Dwb",
    23: "Dwc",
    24: "Dwd",
    25: "Dfa",
    26: "Dfb",
    27: "Dfc",
    28: "Dfd",
    29: "ET",
    30: "EF",
}

MONTH_NAMES = [
    "January",
    "February",
    "March",
    "April",
    "May",
    "June",
    "July",
    "August",
    "September",
    "October",
    "November",
    "December",
]


def _normalize_fips(value: object, width: int) -> str:
    """Return a zero-padded FIPS code with the requested width."""
    text = str(value).strip()
    digits = "".join(ch for ch in text if ch.isdigit())
    if not digits:
        return ""
    return digits.zfill(width)[-width:]


def _load_counties(counties_geojson: Path) -> gpd.GeoDataFrame:
    """Load county polygons and normalize fields used downstream."""
    if not counties_geojson.exists():
        try:
            print(
                f"County GeoJSON not found at {counties_geojson}. "
                f"Attempting download from {DEFAULT_COUNTIES_GEOJSON_URL}..."
            )
            gdf = gpd.read_file(DEFAULT_COUNTIES_GEOJSON_URL)
            counties_geojson.parent.mkdir(parents=True, exist_ok=True)
            # Cache the downloaded file for subsequent runs.
            gdf.to_file(counties_geojson, driver="GeoJSON")
            print(f"Downloaded and cached county GeoJSON to {counties_geojson}")
        except Exception as exc:
            raise FileNotFoundError(
                f"County GeoJSON not found at {counties_geojson}, and download from "
                f"{DEFAULT_COUNTIES_GEOJSON_URL} failed. Download the file manually "
                "and rerun with --counties-geojson pointing to it."
            ) from exc

    gdf = gpd.read_file(counties_geojson)
    if gdf.crs is None:
        gdf = gdf.set_crs("EPSG:4326")
    else:
        gdf = gdf.to_crs("EPSG:4326")

    feature_id = None
    if "id" in gdf.columns:
        feature_id = gdf["id"]
    elif "GEOID" in gdf.columns:
        feature_id = gdf["GEOID"]
    elif "GEOID10" in gdf.columns:
        feature_id = gdf["GEOID10"]
    elif "fips" in gdf.columns:
        feature_id = gdf["fips"]
    else:
        raise ValueError("Unable to locate county FIPS identifier column in county polygons.")

    gdf["county_fips"] = feature_id.map(lambda value: _normalize_fips(value, 5))
    gdf = gdf[gdf["county_fips"] != ""].copy()

    if "NAME" in gdf.columns:
        gdf["county_name"] = gdf["NAME"].fillna("").astype(str).str.strip()
    elif "name" in gdf.columns:
        gdf["county_name"] = gdf["name"].fillna("").astype(str).str.strip()
    else:
        gdf["county_name"] = gdf["county_fips"].map(lambda value: f"County {value}")

    gdf["state_fips"] = gdf["county_fips"].str.slice(0, 2)
    gdf["state"] = gdf["state_fips"].map(lambda code: STATE_FIPS_TO_ABBR.get(code, f"S{code}"))
    gdf = gdf.sort_values("county_fips").reset_index(drop=True)
    return gdf


def _load_koppen_legend(legend_path: Path | None) -> Dict[int, str]:
    """Load Koppen raster codes, using defaults when no legend exists."""
    if legend_path is None:
        return DEFAULT_KOPPEN_CODE_MAP

    mapping: Dict[int, str] = {}
    for line in legend_path.read_text(encoding="utf-8").splitlines():
        text = line.strip()
        if not text or text.startswith("#"):
            continue
        # Handles patterns like:
        # "1:  Af ..." or "1 = Af" or "1 Af"
        import re

        match = re.match(r"^(\d+)\s*[:=]?\s*([A-Za-z]{2,3})\b", text)
        if not match:
            continue

        key = int(match.group(1))
        value = match.group(2)
        mapping[key] = value

    return mapping if mapping else DEFAULT_KOPPEN_CODE_MAP


def _select_data_var(dataset: xr.Dataset, preferred: str) -> str:
    """Choose the best matching climate variable from a dataset."""
    if preferred in dataset.data_vars:
        return preferred

    alias_map = {
        "mlytavg_norm": ["tavg", "tavg_norm"],
        "mlyprcp_norm": ["prcp", "prcp_norm"],
        "dlytmax_norm": ["tmax", "tmax_norm"],
        "dlytmin_norm": ["tmin", "tmin_norm"],
    }
    for alias in alias_map.get(preferred, []):
        if alias in dataset.data_vars:
            return alias

    for candidate in dataset.data_vars:
        if candidate.endswith("_norm"):
            return candidate

    # If the dataset only has one variable, use it as a final fallback.
    if len(dataset.data_vars) == 1:
        return next(iter(dataset.data_vars))

    raise ValueError(
        f"Unable to select a climate variable. Preferred='{preferred}', available={list(dataset.data_vars)}"
    )


def _load_solar_ghi_csv(solar_ghi_csv: Path, counties: gpd.GeoDataFrame) -> List[float]:
    """Load county-keyed average daily GHI values from a representative-point or area-average CSV."""
    with solar_ghi_csv.open(newline="", encoding="utf-8") as handle:
        reader = csv.DictReader(handle)
        if reader.fieldnames is None:
            raise ValueError(f"Solar GHI CSV at {solar_ghi_csv} has no header row.")

        if "county_fips" in reader.fieldnames:
            fips_field = "county_fips"
        elif "countyFips" in reader.fieldnames:
            fips_field = "countyFips"
        else:
            raise ValueError(
                f"Solar GHI CSV at {solar_ghi_csv} must include county_fips or countyFips."
            )

        if "avgSolarGhiKwhM2Day" not in reader.fieldnames:
            raise ValueError(
                f"Solar GHI CSV at {solar_ghi_csv} must include avgSolarGhiKwhM2Day."
            )

        solar_by_fips: Dict[str, float] = {}
        for row in reader:
            county_fips = _normalize_fips(row.get(fips_field, ""), 5)
            raw_value = str(row.get("avgSolarGhiKwhM2Day", "")).strip()
            if not county_fips or not raw_value:
                continue
            try:
                solar_by_fips[county_fips] = float(raw_value)
            except ValueError as exc:
                raise ValueError(
                    f"Invalid avgSolarGhiKwhM2Day value for county {county_fips}: {raw_value}"
                ) from exc

    return [
        solar_by_fips.get(str(row["county_fips"]), float("nan"))
        for _, row in counties.iterrows()
    ]


def _as_monthly_climatology(data_array: xr.DataArray, start_year: int, end_year: int) -> xr.DataArray:
    """Return 12 monthly normals from slices or a time series."""
    if "time" not in data_array.dims:
        raise ValueError(f"Expected a time dimension, got dims={data_array.dims}")

    time_size = int(data_array.sizes.get("time", 0))
    if time_size == 12:
        return data_array

    time_index = data_array["time"]
    if not hasattr(time_index, "dt"):
        # Attempt CF decoding when time is numeric with units/calendar attrs.
        try:
            decoded = xr.decode_cf(xr.Dataset({"_v": data_array}))
            data_array = decoded["_v"]
            time_index = data_array["time"]
        except Exception as exc:
            raise ValueError(
                "Time coordinate does not support datetime access for monthly climatology grouping, "
                "and CF decoding failed."
            ) from exc

    if not hasattr(time_index, "dt"):
        raise ValueError("Time coordinate does not support datetime access for monthly climatology grouping.")

    period = data_array.where(
        (time_index.dt.year >= start_year) & (time_index.dt.year <= end_year),
        drop=True,
    )
    if int(period.sizes.get("time", 0)) == 0:
        raise ValueError(f"No monthly data found between {start_year} and {end_year}.")

    monthly = period.groupby("time.month").mean("time", skipna=True)
    if int(monthly.sizes.get("month", 0)) != 12:
        raise ValueError(
            f"Monthly climatology for {start_year}-{end_year} has {monthly.sizes.get('month', 0)} months; expected 12."
        )

    # Standardize to a `time` dimension so downstream code can reuse .isel(time=month_idx).
    monthly = monthly.rename({"month": "time"})
    monthly = monthly.assign_coords(time=np.arange(1, 13))
    return monthly


def _infer_time_resolution_days(data_array: xr.DataArray) -> float:
    """Estimate the median spacing between time steps in days."""
    time_values = np.asarray(data_array["time"].values)
    if time_values.size < 2:
        return float("inf")

    deltas = np.diff(time_values).astype("timedelta64[D]").astype(np.int64)
    deltas = deltas[deltas > 0]
    if deltas.size == 0:
        return float("inf")
    return float(np.median(deltas))


def _extract_grid_2d(data_array: xr.DataArray) -> Tuple[np.ndarray, "Affine"]:
    """Convert a lat/lon slice to a raster array and transform."""
    # Expected shape for 2D arrays: lat, lon
    # Build affine from center coordinates.
    lon = np.asarray(data_array["lon"].values, dtype=np.float64)
    lat = np.asarray(data_array["lat"].values, dtype=np.float64)
    arr = np.asarray(data_array.values, dtype=np.float64)

    if arr.shape != (lat.size, lon.size):
        raise ValueError(f"Unexpected grid shape {arr.shape}; expected {(lat.size, lon.size)}")

    # Make sure the first row is northernmost to match affine from top-left.
    if lat[0] < lat[-1]:
        lat = lat[::-1]
        arr = arr[::-1, :]

    x_res = abs(lon[1] - lon[0])
    y_res = abs(lat[0] - lat[1])

    from affine import Affine

    top_left_x = lon.min() - (x_res / 2.0)
    top_left_y = lat.max() + (y_res / 2.0)
    transform = Affine.translation(top_left_x, top_left_y) * Affine.scale(x_res, -y_res)
    return arr, transform


def _zonal_mean(arr: np.ndarray, transform, counties: gpd.GeoDataFrame) -> List[float]:
    """Calculate the mean raster value inside each county polygon."""
    values = np.asarray(arr, dtype=np.float64)
    means: List[float] = []
    for geometry in counties.geometry:
        mask = geometry_mask(
            [geometry.__geo_interface__],
            out_shape=values.shape,
            transform=transform,
            invert=True,
            all_touched=True,
        )
        selected = values[mask]
        selected = selected[np.isfinite(selected)]
        means.append(float(selected.mean()) if selected.size else float("nan"))
    return means


def _zonal_mean_raster(raster_path: Path, counties: gpd.GeoDataFrame) -> List[float]:
    """Calculate the mean raster-file value for each county."""
    with rasterio.open(raster_path) as source:
        raster_counties = counties
        if source.crs is not None and counties.crs is not None and counties.crs != source.crs:
            raster_counties = counties.to_crs(source.crs)

        data = source.read(1, masked=True)
        values = np.asarray(data.filled(np.nan), dtype=np.float64)
        if source.nodata is not None:
            values = np.where(values == source.nodata, np.nan, values)
        return _zonal_mean(values, source.transform, raster_counties)


def _zonal_majority_class(koppen_raster: Path, counties: gpd.GeoDataFrame, code_map: Dict[int, str]) -> List[str]:
    """Assign each county its most common Koppen-Geiger class."""
    classes: List[str] = []
    with rasterio.open(koppen_raster) as source:
        raster_counties = counties
        if source.crs is not None and counties.crs is not None and counties.crs != source.crs:
            raster_counties = counties.to_crs(source.crs)

        data = source.read(1, masked=True)
        values = np.asarray(data.filled(0))
        if source.nodata is not None:
            values = np.where(values == source.nodata, 0, values)

        for geometry in raster_counties.geometry:
            mask = geometry_mask(
                [geometry.__geo_interface__],
                out_shape=values.shape,
                transform=source.transform,
                invert=True,
                all_touched=True,
            )
            selected = values[mask]
            selected = selected[selected != 0]
            if selected.size == 0:
                classes.append("Cfa")
                continue

            codes, counts = np.unique(selected.astype(np.int64), return_counts=True)
            code_int = int(codes[int(np.argmax(counts))])
            classes.append(code_map.get(code_int, "Cfa"))
    return classes


def _compute_extreme_days(
    tmax_daily: xr.Dataset,
    tmin_daily: xr.Dataset,
    counties: gpd.GeoDataFrame,
    hot_threshold_c: float,
    freeze_threshold_c: float,
) -> List[float]:
    """Count daily hot or freezing days for each county."""
    tmax_var = _select_data_var(tmax_daily, "dlytmax_norm")
    tmin_var = _select_data_var(tmin_daily, "dlytmin_norm")

    tmax = tmax_daily[tmax_var]
    tmin = tmin_daily[tmin_var]

    # If datasets differ slightly in day count, use the overlapping day count.
    day_count = int(min(tmax.sizes["time"], tmin.sizes["time"]))
    totals = np.zeros(len(counties), dtype=np.float64)

    for day in range(day_count):
        tmax_arr, tmax_transform = _extract_grid_2d(tmax.isel(time=day))
        tmin_arr, tmin_transform = _extract_grid_2d(tmin.isel(time=day))

        if tmax_transform != tmin_transform:
            raise ValueError("Daily tmax/tmin grids do not share the same transform.")

        day_tmax = _zonal_mean(tmax_arr, tmax_transform, counties)
        day_tmin = _zonal_mean(tmin_arr, tmin_transform, counties)

        for idx, (mx, mn) in enumerate(zip(day_tmax, day_tmin)):
            if np.isnan(mx) or np.isnan(mn):
                continue
            if mx >= hot_threshold_c or mn <= freeze_threshold_c:
                totals[idx] += 1.0

    return totals.tolist()


def _compute_extreme_days_monthly_proxy(
    tmax_monthly: xr.DataArray,
    tmin_monthly: xr.DataArray,
    counties: gpd.GeoDataFrame,
    hot_threshold_c: float,
    freeze_threshold_c: float,
) -> List[float]:
    """Estimate extreme days from monthly tmax and tmin means."""
    if int(tmax_monthly.sizes.get("time", 0)) != 12 or int(tmin_monthly.sizes.get("time", 0)) != 12:
        raise ValueError("Monthly proxy for extremeDays requires 12 monthly slices for tmax and tmin.")

    month_days = np.array([31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31], dtype=np.float64)
    totals = np.zeros(len(counties), dtype=np.float64)

    for month in range(12):
        tmax_arr, tmax_transform = _extract_grid_2d(tmax_monthly.isel(time=month))
        tmin_arr, tmin_transform = _extract_grid_2d(tmin_monthly.isel(time=month))
        if tmax_transform != tmin_transform:
            raise ValueError("Monthly proxy tmax/tmin grids do not share the same transform.")

        month_tmax = _zonal_mean(tmax_arr, tmax_transform, counties)
        month_tmin = _zonal_mean(tmin_arr, tmin_transform, counties)

        for idx, (mx, mn) in enumerate(zip(month_tmax, month_tmin)):
            if np.isnan(mx) or np.isnan(mn):
                continue
            if mx >= hot_threshold_c or mn <= freeze_threshold_c:
                totals[idx] += month_days[month]

    return totals.tolist()


def _precip_month_extremes(prcp_months_mm: np.ndarray) -> tuple[str | None, str | None]:
    """Return wettest and driest month names from 12 monthly precipitation totals."""
    valid_mask = np.isfinite(prcp_months_mm)
    if not np.any(valid_mask):
        return None, None

    comparable = np.where(valid_mask, prcp_months_mm, np.nan)
    wettest_month = MONTH_NAMES[int(np.nanargmax(comparable))]
    driest_month = MONTH_NAMES[int(np.nanargmin(comparable))]
    return wettest_month, driest_month


def build_county_records(
    counties_geojson: Path,
    koppen_raster: Path,
    koppen_legend: Path | None,
    monthly_tavg_nc: Path,
    monthly_prcp_nc: Path,
    daily_tmax_nc: Path,
    daily_tmin_nc: Path,
    hot_threshold_f: float,
    freeze_threshold_f: float,
    climatology_start_year: int,
    climatology_end_year: int,
    extreme_days_mode: str,
    solar_ghi_raster: Path | None,
    solar_ghi_csv: Path | None,
) -> Dict[str, dict]:
    """Build county climate records consumed by the web app."""
    counties = _load_counties(counties_geojson)

    koppen_classes = _zonal_majority_class(koppen_raster, counties, _load_koppen_legend(koppen_legend))

    monthly_tavg = xr.open_dataset(monthly_tavg_nc, decode_times=True)
    monthly_prcp = xr.open_dataset(monthly_prcp_nc, decode_times=True)
    daily_tmax = xr.open_dataset(daily_tmax_nc, decode_times=True)
    daily_tmin = xr.open_dataset(daily_tmin_nc, decode_times=True)

    tavg_var = _select_data_var(monthly_tavg, "mlytavg_norm")
    prcp_var = _select_data_var(monthly_prcp, "mlyprcp_norm")

    tavg_monthly = _as_monthly_climatology(
        monthly_tavg[tavg_var],
        start_year=climatology_start_year,
        end_year=climatology_end_year,
    )
    prcp_monthly = _as_monthly_climatology(
        monthly_prcp[prcp_var],
        start_year=climatology_start_year,
        end_year=climatology_end_year,
    )

    tavg_by_month: List[List[float]] = []
    prcp_by_month_mm: List[List[float]] = []
    for month in range(12):
        tavg_arr, tavg_transform = _extract_grid_2d(tavg_monthly.isel(time=month))
        prcp_arr, prcp_transform = _extract_grid_2d(prcp_monthly.isel(time=month))

        if tavg_transform != prcp_transform:
            raise ValueError("Monthly tavg/prcp grids do not share the same transform.")

        tavg_by_month.append(_zonal_mean(tavg_arr, tavg_transform, counties))
        prcp_by_month_mm.append(_zonal_mean(prcp_arr, prcp_transform, counties))

    hot_threshold_c = (hot_threshold_f - 32.0) * (5.0 / 9.0)
    freeze_threshold_c = (freeze_threshold_f - 32.0) * (5.0 / 9.0)

    tmax_var = _select_data_var(daily_tmax, "dlytmax_norm")
    tmin_var = _select_data_var(daily_tmin, "dlytmin_norm")
    tmax_da = daily_tmax[tmax_var]
    tmin_da = daily_tmin[tmin_var]
    resolution_days = min(_infer_time_resolution_days(tmax_da), _infer_time_resolution_days(tmin_da))

    # Daily-like input (true daily normals or daily grids).
    if resolution_days <= 2.0:
        extreme_days = _compute_extreme_days(
            daily_tmax=daily_tmax,
            daily_tmin=daily_tmin,
            counties=counties,
            hot_threshold_c=hot_threshold_c,
            freeze_threshold_c=freeze_threshold_c,
        )
        extreme_days_source_tag = "daily"
    else:
        if extreme_days_mode == "require-daily":
            raise ValueError(
                "Extreme-days inputs appear to be monthly series, but --extreme-days-mode=require-daily was set."
            )

        tmax_monthly = _as_monthly_climatology(
            tmax_da,
            start_year=climatology_start_year,
            end_year=climatology_end_year,
        )
        tmin_monthly = _as_monthly_climatology(
            tmin_da,
            start_year=climatology_start_year,
            end_year=climatology_end_year,
        )
        extreme_days = _compute_extreme_days_monthly_proxy(
            tmax_monthly=tmax_monthly,
            tmin_monthly=tmin_monthly,
            counties=counties,
            hot_threshold_c=hot_threshold_c,
            freeze_threshold_c=freeze_threshold_c,
        )
        extreme_days_source_tag = "monthly-proxy"

    if solar_ghi_raster is not None and solar_ghi_raster.exists():
        solar_ghi_kwh_m2_day = _zonal_mean_raster(solar_ghi_raster, counties)
        solar_source_tag = "solar-ghi-raster"
    elif solar_ghi_csv is not None and solar_ghi_csv.exists():
        solar_ghi_kwh_m2_day = _load_solar_ghi_csv(solar_ghi_csv, counties)
        solar_source_tag = "solar-ghi-representative-point"
    else:
        if solar_ghi_raster is not None:
            print(f"Solar GHI raster not found at {solar_ghi_raster}; leaving avgSolarGhiKwhM2Day blank.")
        if solar_ghi_csv is not None:
            print(f"Solar GHI CSV not found at {solar_ghi_csv}; leaving avgSolarGhiKwhM2Day blank.")
        solar_ghi_kwh_m2_day = [float("nan")] * len(counties)
        solar_source_tag = "no-solar-ghi-source"

    records: Dict[str, dict] = {}
    missing_numeric_count = 0
    for idx, row in counties.iterrows():
        county_fips = row["county_fips"]
        county_name = row["county_name"]
        state = row["state"]
        koppen_zone = koppen_classes[idx]

        temp_months_c = np.array([tavg_by_month[m][idx] for m in range(12)], dtype=np.float64)
        prcp_months_mm = np.array([prcp_by_month_mm[m][idx] for m in range(12)], dtype=np.float64)

        # Guard against nodata counties (outside CONUS grids, islands, etc.).
        valid_temp = temp_months_c[np.isfinite(temp_months_c)]
        valid_prcp = prcp_months_mm[np.isfinite(prcp_months_mm)]

        avg_temp_c = np.nanmean(valid_temp) if valid_temp.size else np.nan
        annual_prcp_mm = np.nansum(valid_prcp) if valid_prcp.size else np.nan

        # Precipitation seasonality as coefficient of variation, capped 0-100.
        if valid_prcp.size:
            mean_prcp = float(np.nanmean(valid_prcp))
            std_prcp = float(np.nanstd(valid_prcp))
            seasonality = 0.0 if mean_prcp <= 0 else max(0.0, min(100.0, (std_prcp / mean_prcp) * 100.0))
        else:
            seasonality = np.nan
        wettest_precip_month, driest_precip_month = _precip_month_extremes(prcp_months_mm)

        use_missing_temp = not np.isfinite(avg_temp_c)
        use_missing_prcp = not np.isfinite(annual_prcp_mm)
        use_missing_seasonality = not np.isfinite(seasonality)
        use_missing_base_noaa = use_missing_temp or use_missing_prcp or use_missing_seasonality
        use_missing_extreme = use_missing_base_noaa or not np.isfinite(extreme_days[idx])
        used_any_missing_numeric = (
            use_missing_temp or use_missing_prcp or use_missing_seasonality or use_missing_extreme
        )

        if used_any_missing_numeric:
            missing_numeric_count += 1

        avg_temp_f = (
            round((float(avg_temp_c) * 9.0 / 5.0) + 32.0, 1)
            if not use_missing_temp
            else None
        )
        annual_prcp_in = (
            round(float(annual_prcp_mm) / 25.4, 1)
            if not use_missing_prcp
            else None
        )
        seasonality_idx = (
            int(round(float(seasonality)))
            if not use_missing_seasonality
            else None
        )
        extreme_days_value = (
            int(round(float(extreme_days[idx])))
            if not use_missing_extreme
            else None
        )
        avg_solar_ghi_value = (
            round(float(solar_ghi_kwh_m2_day[idx]), 2)
            if np.isfinite(solar_ghi_kwh_m2_day[idx])
            else None
        )

        source_suffix = " + missing-noaa-numeric" if used_any_missing_numeric else ""
        solar_source_suffix = "" if avg_solar_ghi_value is not None else " + missing-solar-ghi"

        record = {
            "countyName": county_name,
            "state": state,
            "koppenZone": koppen_zone,
            "avgTempF": avg_temp_f,
            "annualPrecipIn": annual_prcp_in,
            "seasonalityIndex": seasonality_idx,
            "wettestPrecipMonth": wettest_precip_month,
            "driestPrecipMonth": driest_precip_month,
            "extremeDays": extreme_days_value,
            "avgSolarGhiKwhM2Day": avg_solar_ghi_value,
            "source": (
                "kg-beck2023 + noaa-nclimgrid-1991-2020 "
                f"({extreme_days_source_tag}) + {solar_source_tag}{source_suffix}{solar_source_suffix}"
            )
        }
        records[county_fips] = record

    monthly_tavg.close()
    monthly_prcp.close()
    daily_tmax.close()
    daily_tmin.close()
    if missing_numeric_count:
        print(
            f"Marked missing numeric NOAA values for {missing_numeric_count} counties "
            "(likely outside NOAA grid coverage)."
        )
    return records


def write_csv(records: Dict[str, dict], out_file: Path) -> None:
    """Write county records to the browser-loaded CSV payload."""
    fields = [
        "countyFips",
        "countyName",
        "state",
        "koppenZone",
        "avgTempF",
        "annualPrecipIn",
        "seasonalityIndex",
        "wettestPrecipMonth",
        "driestPrecipMonth",
        "extremeDays",
        "avgSolarGhiKwhM2Day",
        "source",
    ]

    with out_file.open("w", encoding="utf-8", newline="") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        for county_fips in sorted(records):
            row = {"countyFips": county_fips}
            row.update(records[county_fips])
            writer.writerow(row)


def parse_args() -> argparse.Namespace:
    """Define and parse command-line options for this generator."""
    parser = argparse.ArgumentParser(description="Generate county climate records for the web app.")
    parser.add_argument("--counties-geojson", type=Path, required=True, help="County polygon GeoJSON path.")
    parser.add_argument("--koppen-raster", type=Path, required=True, help="Koppen-Geiger raster TIFF path.")
    parser.add_argument("--koppen-legend", type=Path, default=None, help="Optional legend.txt mapping integer codes.")
    parser.add_argument(
        "--monthly-tavg-nc",
        type=Path,
        required=True,
        help="NOAA monthly tavg netCDF (12-slice normals or monthly time-series).",
    )
    parser.add_argument(
        "--monthly-prcp-nc",
        type=Path,
        required=True,
        help="NOAA monthly prcp netCDF (12-slice normals or monthly time-series).",
    )
    parser.add_argument(
        "--daily-tmax-nc",
        type=Path,
        required=True,
        help="NOAA tmax netCDF (daily normals/grids preferred; monthly time-series allowed in proxy mode).",
    )
    parser.add_argument(
        "--daily-tmin-nc",
        type=Path,
        required=True,
        help="NOAA tmin netCDF (daily normals/grids preferred; monthly time-series allowed in proxy mode).",
    )
    parser.add_argument("--hot-threshold-f", type=float, default=95.0, help="Hot day threshold in Fahrenheit.")
    parser.add_argument("--freeze-threshold-f", type=float, default=32.0, help="Freeze day threshold in Fahrenheit.")
    parser.add_argument(
        "--climatology-start-year",
        type=int,
        default=1991,
        help="Start year (inclusive) for monthly climatology calculation when time series files are provided.",
    )
    parser.add_argument(
        "--climatology-end-year",
        type=int,
        default=2020,
        help="End year (inclusive) for monthly climatology calculation when time series files are provided.",
    )
    parser.add_argument(
        "--extreme-days-mode",
        choices=["auto", "require-daily"],
        default="auto",
        help="`auto` allows monthly-proxy extremeDays if daily grids are not provided; `require-daily` enforces daily input.",
    )
    parser.add_argument(
        "--solar-ghi-raster",
        type=Path,
        default=None,
        help="Optional raster of annual average daily GHI in kWh/m2/day for avgSolarGhiKwhM2Day.",
    )
    parser.add_argument(
        "--solar-ghi-csv",
        type=Path,
        default=None,
        help=(
            "Optional county CSV with avgSolarGhiKwhM2Day. Used as a representative-point fallback "
            "when --solar-ghi-raster is not supplied."
        ),
    )
    parser.add_argument("--out", type=Path, default=Path("data/climate-data.csv"), help="Output CSV file path.")
    return parser.parse_args()


def main() -> None:
    """Run the ETL flow and write the final CSV file."""
    args = parse_args()
    records = build_county_records(
        counties_geojson=args.counties_geojson,
        koppen_raster=args.koppen_raster,
        koppen_legend=args.koppen_legend,
        monthly_tavg_nc=args.monthly_tavg_nc,
        monthly_prcp_nc=args.monthly_prcp_nc,
        daily_tmax_nc=args.daily_tmax_nc,
        daily_tmin_nc=args.daily_tmin_nc,
        hot_threshold_f=args.hot_threshold_f,
        freeze_threshold_f=args.freeze_threshold_f,
        climatology_start_year=args.climatology_start_year,
        climatology_end_year=args.climatology_end_year,
        extreme_days_mode=args.extreme_days_mode,
        solar_ghi_raster=args.solar_ghi_raster,
        solar_ghi_csv=args.solar_ghi_csv,
    )

    write_csv(records=records, out_file=args.out)
    print(f"Wrote {len(records)} county records to {args.out}")


if __name__ == "__main__":
    main()
