#!/usr/bin/env python3
"""
Download NOAA nClimGrid-Daily county temperature files and build locally
extreme days/year plus a daily absolute extreme climate bucket.

The metric uses county-specific historical thresholds:

- p95_tmax_c: 95th percentile of county-average daily Tmax during 1991-2020
- p05_tmin_c: 5th percentile of county-average daily Tmin during 1991-2020
- locallyExtremeDays: count of days in a year where
  Tmax >= p95_tmax_c OR Tmin <= p05_tmin_c
- absoluteExtremeDays: count of days in a year where
  Tmax >= 95F OR Tmin <= 0F by default

NOAA source files are cached under:

data/noaa/tmax_cty_scaled/YYYY/tmax-YYYYMM-cty-scaled.csv
data/noaa/tmin_cty_scaled/YYYY/tmin-YYYYMM-cty-scaled.csv
"""

from __future__ import annotations

import argparse
import calendar
import csv
import math
import shutil
import time
import urllib.error
import urllib.request
from array import array
from collections import defaultdict
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Dict, Iterable, Iterator, List, Tuple


NOAA_AVERAGES_BASE_URL = "https://www.ncei.noaa.gov/data/nclimgrid-daily/access/averages"
NOAA_STATE_CROSSWALK_URL = (
    "https://www.ncei.noaa.gov/data/nclimgrid-daily/doc/us-state-codes_ncei-to-fips.csv"
)

REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_NOAA_DIR = REPO_ROOT / "data" / "noaa"
DEFAULT_TMAX_DIR = DEFAULT_NOAA_DIR / "tmax_cty_scaled"
DEFAULT_TMIN_DIR = DEFAULT_NOAA_DIR / "tmin_cty_scaled"
DEFAULT_METADATA_DIR = DEFAULT_NOAA_DIR / "nclimgrid"
DEFAULT_STATE_CROSSWALK = DEFAULT_METADATA_DIR / "us-state-codes_ncei-to-fips.csv"
DEFAULT_OUT = DEFAULT_NOAA_DIR / "county_locally_extreme_days.csv"
DEFAULT_THRESHOLDS_OUT = DEFAULT_NOAA_DIR / "county_locally_extreme_thresholds.csv"
DEFAULT_COMPARISON_OUT = DEFAULT_NOAA_DIR / "county_locally_extreme_days_comparison.csv"
DEFAULT_OLD_CLIMATE_DATA = REPO_ROOT / "data" / "climate-data.csv"

MISSING_VALUE = -999.99

# NOAA carries the District of Columbia row with an NCEI-style region code that
# cannot be translated by the state-code crosswalk alone.
NOAA_REGION_CODE_TO_FIPS_OVERRIDES = {
    "18511": "11001",  # DC: District of Columbia
}

# One-to-one county-equivalent vintage changes that are safe enough for
# comparison joins. Splits are intentionally excluded.
OLD_APP_FIPS_TO_CURRENT_FIPS = {
    "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.",
    ),
    "02270": (
        "02158",
        "Wade Hampton Census Area, AK changed name/code to Kusilvak Census Area effective 2015-07-01.",
    ),
}


@dataclass(frozen=True)
class MonthlyCountyValues:
    """One county row from a NOAA county area-average monthly CSV."""

    region_code: str
    county_fips: str
    county_name: str
    values: Tuple[float | None, ...]


@dataclass
class AnnualCounts:
    """Annual hot/cold/combined locally extreme day counts for one county."""

    hot: int = 0
    cold: int = 0
    combined: int = 0
    absolute_combined: int = 0
    valid_days: int = 0


def c_to_f(value_c: float) -> float:
    return (value_c * 9.0 / 5.0) + 32.0


def f_to_c(value_f: float) -> float:
    return (value_f - 32.0) * (5.0 / 9.0)


def percentile(values: Iterable[float], pct: float) -> float:
    """Return a linear-interpolated percentile, matching numpy's default style."""
    sorted_values = sorted(values)
    if not sorted_values:
        return math.nan

    if len(sorted_values) == 1:
        return float(sorted_values[0])

    position = (len(sorted_values) - 1) * (pct / 100.0)
    lower_index = int(math.floor(position))
    upper_index = int(math.ceil(position))
    if lower_index == upper_index:
        return float(sorted_values[lower_index])

    fraction = position - lower_index
    lower = sorted_values[lower_index]
    upper = sorted_values[upper_index]
    return float(lower + ((upper - lower) * fraction))


def canonical_noaa_filename(variable: str, year: int, month: int) -> str:
    return f"{variable}-{year}{month:02d}-cty-scaled.csv"


def canonical_noaa_url(variable: str, year: int, month: int) -> str:
    return f"{NOAA_AVERAGES_BASE_URL}/{year}/{canonical_noaa_filename(variable, year, month)}"


def canonical_cache_path(variable_dir: Path, variable: str, year: int, month: int) -> Path:
    return variable_dir / str(year) / canonical_noaa_filename(variable, year, month)


def candidate_cache_paths(variable_dir: Path, variable: str, year: int, month: int) -> List[Path]:
    """Support the canonical filename plus simple month-file placeholders."""
    year_dir = variable_dir / str(year)
    return [
        year_dir / canonical_noaa_filename(variable, year, month),
        year_dir / f"{month:02d}.csv",
        year_dir / f"{month:02d}",
    ]


def resolve_cache_path(variable_dir: Path, variable: str, year: int, month: int) -> Path:
    """Return the first usable cached NOAA file path, falling back to canonical."""
    for path in candidate_cache_paths(variable_dir, variable, year, month):
        if path.exists() and path.is_file() and path.stat().st_size > 0:
            return path
    return canonical_cache_path(variable_dir, variable, year, month)


def download_file(url: str, target: Path, force: bool, timeout: int, retries: int) -> bool:
    """Download a URL to target. Returns True when a new file was written."""
    if target.exists() and target.stat().st_size > 0 and not force:
        return False

    target.parent.mkdir(parents=True, exist_ok=True)
    last_error: Exception | None = None

    for attempt in range(1, retries + 1):
        try:
            request = urllib.request.Request(
                url,
                headers={
                    "User-Agent": (
                        "Temperature-Based-Analysis county locally extreme days builder "
                        "(contact: local research script)"
                    )
                },
            )
            with urllib.request.urlopen(request, timeout=timeout) as response:
                with NamedTemporaryFile("wb", delete=False, dir=target.parent) as tmp_file:
                    shutil.copyfileobj(response, tmp_file)
                    temp_path = Path(tmp_file.name)

            if temp_path.stat().st_size == 0:
                temp_path.unlink(missing_ok=True)
                raise RuntimeError(f"Downloaded empty file from {url}")

            temp_path.replace(target)
            return True
        except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, RuntimeError) as exc:
            last_error = exc
            if attempt < retries:
                time.sleep(min(10.0, 1.5 * attempt))

    raise RuntimeError(f"Failed to download {url} after {retries} attempts: {last_error}")


def download_noaa_temperature_files(
    *,
    start_year: int,
    end_year: int,
    tmax_dir: Path,
    tmin_dir: Path,
    force: bool,
    timeout: int,
    retries: int,
    sleep_seconds: float,
) -> None:
    """Cache all monthly county Tmax/Tmin files needed for the metric."""
    total_written = 0
    total_skipped = 0

    for year in range(start_year, end_year + 1):
        for month in range(1, 13):
            for variable, variable_dir in (("tmax", tmax_dir), ("tmin", tmin_dir)):
                target = canonical_cache_path(variable_dir, variable, year, month)
                url = canonical_noaa_url(variable, year, month)
                written = download_file(
                    url=url,
                    target=target,
                    force=force,
                    timeout=timeout,
                    retries=retries,
                )
                if written:
                    total_written += 1
                    print(f"Downloaded {target}")
                    if sleep_seconds > 0:
                        time.sleep(sleep_seconds)
                else:
                    total_skipped += 1

    print(f"NOAA cache ready: {total_written} downloaded, {total_skipped} already present.")


def download_state_crosswalk(path: Path, *, force: bool, timeout: int, retries: int) -> None:
    """Cache the NCEI state-code to FIPS state-code crosswalk."""
    written = download_file(
        url=NOAA_STATE_CROSSWALK_URL,
        target=path,
        force=force,
        timeout=timeout,
        retries=retries,
    )
    if written:
        print(f"Downloaded NCEI/FIPS state crosswalk to {path}")


def _numeric_code(text: object, width: int) -> str:
    """Extract and zero-pad a fixed-width numeric code."""
    digits = "".join(ch for ch in str(text).strip() if ch.isdigit())
    return digits.zfill(width)[-width:] if digits else ""


def load_state_crosswalk(path: Path) -> Dict[str, str]:
    """Load NCEI state code -> FIPS state code mapping."""
    if not path.exists():
        print(
            f"State crosswalk not found at {path}. Falling back to the first two "
            "county 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,
    )

    mapping: Dict[str, str] = {}
    data_rows = rows[1:] if ncei_idx is not None and fips_idx is not None else rows

    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 = _numeric_code(row[ncei_idx], 2)
            fips = _numeric_code(row[fips_idx], 2)
        else:
            numeric_cells = [_numeric_code(cell, 2) for cell in row]
            numeric_cells = [cell for cell in numeric_cells if cell]
            if len(numeric_cells) < 2:
                continue
            ncei, fips = numeric_cells[0], numeric_cells[1]

        if ncei and fips:
            mapping[ncei] = fips

    print(f"Loaded {len(mapping)} NCEI state-code mappings from {path}")
    return mapping


def county_fips_from_region_code(region_code: str, state_crosswalk: Dict[str, str]) -> str:
    """Convert a NOAA county region code to a current 5-digit county FIPS."""
    code = _numeric_code(region_code, 5)
    if len(code) != 5:
        return ""

    if code in NOAA_REGION_CODE_TO_FIPS_OVERRIDES:
        return NOAA_REGION_CODE_TO_FIPS_OVERRIDES[code]

    ncei_state = code[:2]
    county_code = code[-3:]
    state_fips = state_crosswalk.get(ncei_state, ncei_state)
    return f"{state_fips}{county_code}" if state_fips else ""


def iter_monthly_county_values(
    path: Path,
    *,
    year: int,
    month: int,
    state_crosswalk: Dict[str, str],
) -> Iterator[MonthlyCountyValues]:
    """Yield county daily values from one NOAA area-average CSV."""
    days_in_month = calendar.monthrange(year, month)[1]

    with path.open("r", encoding="utf-8-sig", newline="") as csv_file:
        reader = csv.reader(csv_file)
        for row in reader:
            if len(row) < 6:
                continue

            region_type = row[0].strip().lower()
            if region_type in {"region type", "region_type"}:
                continue
            if region_type != "cty":
                continue

            region_code = _numeric_code(row[1], 5)
            county_fips = county_fips_from_region_code(region_code, state_crosswalk)
            county_name = row[2].strip()
            raw_values = row[6 : 6 + days_in_month]

            values: List[float | None] = []
            for raw_value in raw_values:
                try:
                    value = float(raw_value)
                except ValueError:
                    values.append(None)
                    continue

                if value <= MISSING_VALUE + 0.001:
                    values.append(None)
                else:
                    values.append(value)

            yield MonthlyCountyValues(
                region_code=region_code,
                county_fips=county_fips,
                county_name=county_name,
                values=tuple(values),
            )


def require_monthly_file(variable_dir: Path, variable: str, year: int, month: int) -> Path:
    """Return a cached monthly NOAA file or raise a specific missing-file error."""
    path = resolve_cache_path(variable_dir, variable, year, month)
    if path.exists() and path.is_file() and path.stat().st_size > 0:
        return path

    raise FileNotFoundError(
        f"Missing {variable} county file for {year}-{month:02d}. Expected one of: "
        + ", ".join(str(candidate) for candidate in candidate_cache_paths(variable_dir, variable, year, month))
    )


def build_thresholds(
    *,
    baseline_start_year: int,
    baseline_end_year: int,
    tmax_dir: Path,
    tmin_dir: Path,
    state_crosswalk: Dict[str, str],
) -> Tuple[Dict[str, dict], Dict[str, dict]]:
    """Build per-county baseline percentile thresholds."""
    tmax_values: Dict[str, array] = {}
    tmin_values: Dict[str, array] = {}
    county_meta: Dict[str, dict] = {}

    for year in range(baseline_start_year, baseline_end_year + 1):
        print(f"Reading baseline year {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)

            for record in iter_monthly_county_values(
                tmax_path,
                year=year,
                month=month,
                state_crosswalk=state_crosswalk,
            ):
                if not record.county_fips:
                    continue
                county_meta.setdefault(
                    record.county_fips,
                    {
                        "region_code": record.region_code,
                        "county_name": record.county_name,
                    },
                )
                county_values = tmax_values.setdefault(record.county_fips, array("f"))
                county_values.extend(value for value in record.values if value is not None)

            for record in iter_monthly_county_values(
                tmin_path,
                year=year,
                month=month,
                state_crosswalk=state_crosswalk,
            ):
                if not record.county_fips:
                    continue
                county_meta.setdefault(
                    record.county_fips,
                    {
                        "region_code": record.region_code,
                        "county_name": record.county_name,
                    },
                )
                county_values = tmin_values.setdefault(record.county_fips, array("f"))
                county_values.extend(value for value in record.values if value is not None)

    thresholds: Dict[str, dict] = {}
    for county_fips in sorted(county_meta):
        county_tmax = tmax_values.get(county_fips, array("f"))
        county_tmin = tmin_values.get(county_fips, array("f"))
        thresholds[county_fips] = {
            "p95_tmax_c": percentile(county_tmax, 95.0),
            "p05_tmin_c": percentile(county_tmin, 5.0),
            "baseline_tmax_days": len(county_tmax),
            "baseline_tmin_days": len(county_tmin),
        }

    return thresholds, county_meta


def rows_by_fips(
    path: Path,
    *,
    year: int,
    month: int,
    state_crosswalk: Dict[str, str],
) -> Dict[str, MonthlyCountyValues]:
    """Load one monthly NOAA CSV into county rows keyed by county FIPS."""
    return {
        record.county_fips: record
        for record in iter_monthly_county_values(
            path,
            year=year,
            month=month,
            state_crosswalk=state_crosswalk,
        )
        if record.county_fips
    }


def build_annual_counts(
    *,
    analysis_start_year: int,
    analysis_end_year: int,
    tmax_dir: Path,
    tmin_dir: Path,
    state_crosswalk: Dict[str, str],
    thresholds: Dict[str, dict],
    county_meta: Dict[str, dict],
    absolute_heat_threshold_c: float,
    absolute_cold_threshold_c: float,
) -> Dict[Tuple[str, int], AnnualCounts]:
    """Count local and absolute extreme days for each county/year."""
    annual_counts: Dict[Tuple[str, int], AnnualCounts] = {}

    for year in range(analysis_start_year, analysis_end_year + 1):
        print(f"Counting locally extreme days 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)):
                threshold = thresholds.get(county_fips)
                if not threshold:
                    continue

                p95_tmax_c = threshold["p95_tmax_c"]
                p05_tmin_c = threshold["p05_tmin_c"]
                if not math.isfinite(p95_tmax_c) or not math.isfinite(p05_tmin_c):
                    continue

                tmax_record = tmax_rows[county_fips]
                tmin_record = tmin_rows[county_fips]
                county_meta.setdefault(
                    county_fips,
                    {
                        "region_code": tmax_record.region_code,
                        "county_name": tmax_record.county_name,
                    },
                )
                counts = annual_counts.setdefault((county_fips, year), AnnualCounts())

                for tmax_value, tmin_value in zip(tmax_record.values, tmin_record.values):
                    if tmax_value is None or tmin_value is None:
                        continue

                    hot = tmax_value >= p95_tmax_c
                    cold = tmin_value <= p05_tmin_c
                    absolute_heat = tmax_value >= absolute_heat_threshold_c
                    absolute_cold = tmin_value <= absolute_cold_threshold_c
                    counts.valid_days += 1
                    if hot:
                        counts.hot += 1
                    if cold:
                        counts.cold += 1
                    if hot or cold:
                        counts.combined += 1
                    if absolute_heat or absolute_cold:
                        counts.absolute_combined += 1

    return annual_counts


def write_thresholds_csv(
    *,
    path: Path,
    thresholds: Dict[str, dict],
    county_meta: Dict[str, dict],
    baseline_start_year: int,
    baseline_end_year: int,
) -> None:
    """Write the per-county baseline percentile thresholds CSV."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "countyFips",
        "noaaRegionCode",
        "countyName",
        "p95TmaxC",
        "p95TmaxF",
        "p05TminC",
        "p05TminF",
        "baselineValidTmaxDays",
        "baselineValidTminDays",
        "baselineStartYear",
        "baselineEndYear",
        "source",
    ]

    with path.open("w", encoding="utf-8", newline="") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fields)
        writer.writeheader()
        for county_fips in sorted(thresholds):
            threshold = thresholds[county_fips]
            meta = county_meta.get(county_fips, {})
            p95_tmax_c = threshold["p95_tmax_c"]
            p05_tmin_c = threshold["p05_tmin_c"]
            writer.writerow(
                {
                    "countyFips": county_fips,
                    "noaaRegionCode": meta.get("region_code", ""),
                    "countyName": meta.get("county_name", ""),
                    "p95TmaxC": round(p95_tmax_c, 2) if math.isfinite(p95_tmax_c) else "",
                    "p95TmaxF": round(c_to_f(p95_tmax_c), 2) if math.isfinite(p95_tmax_c) else "",
                    "p05TminC": round(p05_tmin_c, 2) if math.isfinite(p05_tmin_c) else "",
                    "p05TminF": round(c_to_f(p05_tmin_c), 2) if math.isfinite(p05_tmin_c) else "",
                    "baselineValidTmaxDays": threshold["baseline_tmax_days"],
                    "baselineValidTminDays": threshold["baseline_tmin_days"],
                    "baselineStartYear": baseline_start_year,
                    "baselineEndYear": baseline_end_year,
                    "source": "noaa-nclimgrid-daily-county-area-averages-scaled",
                }
            )

    print(f"Wrote county percentile thresholds to {path}")


def write_annual_counts_csv(
    *,
    path: Path,
    annual_counts: Dict[Tuple[str, int], AnnualCounts],
    thresholds: Dict[str, dict],
    county_meta: Dict[str, dict],
    baseline_start_year: int,
    baseline_end_year: int,
    absolute_heat_threshold_f: float,
    absolute_cold_threshold_f: float,
) -> None:
    """Write annual county-level locally and absolutely extreme day counts."""
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "countyFips",
        "noaaRegionCode",
        "countyName",
        "year",
        "p95TmaxC",
        "p05TminC",
        "hotExtremeDays",
        "coldExtremeDays",
        "locallyExtremeDays",
        "absoluteExtremeDays",
        "validDays",
        "absoluteExtremeHeatThresholdF",
        "absoluteExtremeColdThresholdF",
        "baselineStartYear",
        "baselineEndYear",
        "source",
    ]

    with path.open("w", encoding="utf-8", newline="") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fields)
        writer.writeheader()
        for county_fips, year in sorted(annual_counts):
            counts = annual_counts[(county_fips, year)]
            threshold = thresholds[county_fips]
            meta = county_meta.get(county_fips, {})
            writer.writerow(
                {
                    "countyFips": county_fips,
                    "noaaRegionCode": meta.get("region_code", ""),
                    "countyName": meta.get("county_name", ""),
                    "year": year,
                    "p95TmaxC": round(threshold["p95_tmax_c"], 2),
                    "p05TminC": round(threshold["p05_tmin_c"], 2),
                    "hotExtremeDays": counts.hot,
                    "coldExtremeDays": counts.cold,
                    "locallyExtremeDays": counts.combined,
                    "absoluteExtremeDays": counts.absolute_combined,
                    "validDays": counts.valid_days,
                    "absoluteExtremeHeatThresholdF": absolute_heat_threshold_f,
                    "absoluteExtremeColdThresholdF": absolute_cold_threshold_f,
                    "baselineStartYear": baseline_start_year,
                    "baselineEndYear": baseline_end_year,
                    "source": "noaa-nclimgrid-daily-county-area-averages-scaled",
                }
            )

    print(f"Wrote locally extreme annual counts to {path}")


def _split_noaa_county_name(name: str) -> Tuple[str, str]:
    """Split NOAA's 'STATE: County Name County' label into app-style fields."""
    if ": " in name:
        state, county_name = name.split(": ", 1)
        return state.strip(), county_name.replace(" County", "").strip()
    return "", name.replace(" County", "").strip()


def _average_or_none(values: Iterable[float]) -> float | None:
    """Average finite values, returning None when no numeric values are present."""
    numbers = [value for value in values if math.isfinite(value)]
    if not numbers:
        return None
    return sum(numbers) / len(numbers)


def _format_number(value: float | None, digits: int = 2) -> str | float:
    """Format optional numeric CSV values as rounded numbers or blanks."""
    if value is None or not math.isfinite(value):
        return ""
    return round(value, digits)


def _current_fips_for_old_app_fips(fips: str) -> Tuple[str, str]:
    """Return the current FIPS and note for known one-to-one county changes."""
    if fips in OLD_APP_FIPS_TO_CURRENT_FIPS:
        return OLD_APP_FIPS_TO_CURRENT_FIPS[fips]
    return fips, ""


def write_comparison_csv(
    *,
    path: Path,
    annual_counts: Dict[Tuple[str, int], AnnualCounts],
    county_meta: Dict[str, dict],
    old_climate_data: Path,
    absolute_heat_threshold_f: float,
    absolute_cold_threshold_f: float,
) -> None:
    """Compare current-vintage daily extreme metrics with the old app metric."""
    if not old_climate_data.exists():
        print(f"Old climate-data CSV not found at {old_climate_data}; skipping comparison CSV.")
        return

    new_by_fips: Dict[str, dict] = {}
    for county_fips, year in sorted(annual_counts):
        row = new_by_fips.setdefault(
            county_fips,
            {
                "years": [],
                "locally": [],
                "hot": [],
                "cold": [],
                "absolute": [],
            },
        )
        counts = annual_counts[(county_fips, year)]
        row["years"].append(year)
        row["locally"].append(counts.combined)
        row["hot"].append(counts.hot)
        row["cold"].append(counts.cold)
        row["absolute"].append(counts.absolute_combined)

    old_by_current_fips: Dict[str, List[dict]] = defaultdict(list)
    with old_climate_data.open("r", encoding="utf-8-sig", newline="") as csv_file:
        reader = csv.DictReader(csv_file)
        for row in reader:
            old_fips = _numeric_code(row.get("countyFips", ""), 5)
            if not old_fips:
                continue
            current_fips, adjustment_note = _current_fips_for_old_app_fips(old_fips)
            extreme_days_raw = (row.get("oldExtremeDays") or row.get("extremeDays") or "").strip()
            try:
                old_extreme_days = float(extreme_days_raw) if extreme_days_raw else None
            except ValueError:
                old_extreme_days = None

            old_by_current_fips[current_fips].append(
                {
                    "old_fips": old_fips,
                    "county_name": row.get("countyName", "").strip(),
                    "state": row.get("state", "").strip(),
                    "old_extreme_days": old_extreme_days,
                    "adjustment_note": adjustment_note,
                }
            )

    path.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "countyFips",
        "noaaRegionCode",
        "countyName",
        "state",
        "hasLocallyExtremeData",
        "hasOldExtremeDays",
        "analysisStartYear",
        "analysisEndYear",
        "analysisYears",
        "avgLocallyExtremeDays",
        "avgHotExtremeDays",
        "avgColdExtremeDays",
        "avgAbsoluteExtremeDays",
        "absoluteExtremeHeatThresholdF",
        "absoluteExtremeColdThresholdF",
        "oldExtremeDays",
        "differenceVsOld",
        "ratioVsOld",
        "oldSourceFips",
        "oldSourceCountyNames",
        "fipsAdjustment",
        "comparisonNote",
    ]

    rows: List[dict] = []
    for county_fips in sorted(set(new_by_fips) | set(old_by_current_fips)):
        new_row = new_by_fips.get(county_fips)
        old_rows = old_by_current_fips.get(county_fips, [])
        old_values = [
            row["old_extreme_days"]
            for row in old_rows
            if row["old_extreme_days"] is not None
        ]
        old_average = _average_or_none(old_values)

        avg_locally = _average_or_none(new_row["locally"]) if new_row else None
        avg_hot = _average_or_none(new_row["hot"]) if new_row else None
        avg_cold = _average_or_none(new_row["cold"]) if new_row else None
        avg_absolute = _average_or_none(new_row["absolute"]) if new_row else None
        diff = avg_locally - old_average if avg_locally is not None and old_average is not None else None
        ratio = (
            avg_locally / old_average
            if avg_locally is not None and old_average not in (None, 0.0)
            else None
        )

        meta = county_meta.get(county_fips, {})
        noaa_state, noaa_county_name = _split_noaa_county_name(str(meta.get("county_name", "")))
        direct_old_rows = [row for row in old_rows if row["old_fips"] == county_fips]
        best_old_row = direct_old_rows[0] if direct_old_rows else (old_rows[0] if old_rows else {})
        county_name = noaa_county_name or best_old_row.get("county_name", "")
        state = noaa_state or best_old_row.get("state", "")

        adjustment_notes = sorted(
            {
                row["adjustment_note"]
                for row in old_rows
                if row.get("adjustment_note")
            }
        )
        has_new = new_row is not None
        has_old_metric = old_average is not None
        if has_new and has_old_metric:
            comparison_note = "matched"
            if adjustment_notes:
                comparison_note = "matched-after-fips-concordance"
        elif has_new:
            comparison_note = "locally-extreme-only; no matching old app metric after FIPS concordance"
        elif has_old_metric:
            comparison_note = "old-metric-only; no NOAA locally extreme county row after FIPS concordance"
        else:
            comparison_note = "old-app-row-without-old-extremeDays-and-no-NOAA-row"

        rows.append(
            {
                "countyFips": county_fips,
                "noaaRegionCode": meta.get("region_code", ""),
                "countyName": county_name,
                "state": state,
                "hasLocallyExtremeData": "yes" if has_new else "no",
                "hasOldExtremeDays": "yes" if has_old_metric else "no",
                "analysisStartYear": min(new_row["years"]) if new_row else "",
                "analysisEndYear": max(new_row["years"]) if new_row else "",
                "analysisYears": len(new_row["years"]) if new_row else "",
                "avgLocallyExtremeDays": _format_number(avg_locally),
                "avgHotExtremeDays": _format_number(avg_hot),
                "avgColdExtremeDays": _format_number(avg_cold),
                "avgAbsoluteExtremeDays": _format_number(avg_absolute),
                "absoluteExtremeHeatThresholdF": absolute_heat_threshold_f if has_new else "",
                "absoluteExtremeColdThresholdF": absolute_cold_threshold_f if has_new else "",
                "oldExtremeDays": _format_number(old_average),
                "differenceVsOld": _format_number(diff),
                "ratioVsOld": _format_number(ratio, digits=3),
                "oldSourceFips": ";".join(row["old_fips"] for row in old_rows),
                "oldSourceCountyNames": ";".join(row["county_name"] for row in old_rows),
                "fipsAdjustment": " | ".join(adjustment_notes),
                "comparisonNote": comparison_note,
            }
        )

    with path.open("w", encoding="utf-8", newline="") as csv_file:
        writer = csv.DictWriter(csv_file, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)

    print(f"Wrote FIPS-aware locally extreme comparison to {path}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Download NOAA nClimGrid-Daily county Tmax/Tmin files and compute "
            "locally extreme days/year."
        )
    )
    parser.add_argument("--baseline-start-year", type=int, default=1991)
    parser.add_argument("--baseline-end-year", type=int, default=2020)
    parser.add_argument(
        "--analysis-start-year",
        type=int,
        default=1991,
        help="First year to count locally extreme days for.",
    )
    parser.add_argument(
        "--analysis-end-year",
        type=int,
        default=date.today().year - 1,
        help="Last year to count locally extreme days for. Defaults to the latest likely complete year.",
    )
    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)
    parser.add_argument("--thresholds-out", type=Path, default=DEFAULT_THRESHOLDS_OUT)
    parser.add_argument("--comparison-out", type=Path, default=DEFAULT_COMPARISON_OUT)
    parser.add_argument(
        "--absolute-heat-threshold-f",
        type=float,
        default=95.0,
        help="Daily Tmax threshold for absoluteExtremeDays, in Fahrenheit.",
    )
    parser.add_argument(
        "--absolute-cold-threshold-f",
        type=float,
        default=0.0,
        help="Daily Tmin threshold for absoluteExtremeDays, in Fahrenheit.",
    )
    parser.add_argument(
        "--old-climate-data",
        type=Path,
        default=DEFAULT_OLD_CLIMATE_DATA,
        help="Existing app climate-data.csv used to compare against the old extremeDays metric.",
    )
    parser.add_argument("--skip-comparison", action="store_true", help="Do not write the comparison CSV.")
    parser.add_argument("--skip-download", action="store_true", help="Use existing local cache only.")
    parser.add_argument("--download-only", action="store_true", help="Download/cache files without computing outputs.")
    parser.add_argument("--force-download", action="store_true", help="Redownload files even if they already exist.")
    parser.add_argument("--timeout", type=int, default=120, help="Per-file download timeout in seconds.")
    parser.add_argument("--retries", type=int, default=3, help="Download attempts per file.")
    parser.add_argument(
        "--sleep-seconds",
        type=float,
        default=0.05,
        help="Pause between newly downloaded NOAA files.",
    )
    return parser.parse_args()


def validate_years(args: argparse.Namespace) -> None:
    """Validate that baseline and analysis year ranges are ordered."""
    if args.baseline_start_year > args.baseline_end_year:
        raise ValueError("--baseline-start-year must be <= --baseline-end-year")
    if args.analysis_start_year > args.analysis_end_year:
        raise ValueError("--analysis-start-year must be <= --analysis-end-year")


def main() -> None:
    args = parse_args()
    validate_years(args)

    download_start_year = min(args.baseline_start_year, args.analysis_start_year)
    download_end_year = max(args.baseline_end_year, args.analysis_end_year)

    if not args.skip_download:
        download_state_crosswalk(
            args.state_crosswalk,
            force=args.force_download,
            timeout=args.timeout,
            retries=args.retries,
        )
        download_noaa_temperature_files(
            start_year=download_start_year,
            end_year=download_end_year,
            tmax_dir=args.tmax_dir,
            tmin_dir=args.tmin_dir,
            force=args.force_download,
            timeout=args.timeout,
            retries=args.retries,
            sleep_seconds=args.sleep_seconds,
        )

    if args.download_only:
        print("Download-only mode complete.")
        return

    state_crosswalk = load_state_crosswalk(args.state_crosswalk)
    thresholds, county_meta = build_thresholds(
        baseline_start_year=args.baseline_start_year,
        baseline_end_year=args.baseline_end_year,
        tmax_dir=args.tmax_dir,
        tmin_dir=args.tmin_dir,
        state_crosswalk=state_crosswalk,
    )
    write_thresholds_csv(
        path=args.thresholds_out,
        thresholds=thresholds,
        county_meta=county_meta,
        baseline_start_year=args.baseline_start_year,
        baseline_end_year=args.baseline_end_year,
    )

    annual_counts = build_annual_counts(
        analysis_start_year=args.analysis_start_year,
        analysis_end_year=args.analysis_end_year,
        tmax_dir=args.tmax_dir,
        tmin_dir=args.tmin_dir,
        state_crosswalk=state_crosswalk,
        thresholds=thresholds,
        county_meta=county_meta,
        absolute_heat_threshold_c=f_to_c(args.absolute_heat_threshold_f),
        absolute_cold_threshold_c=f_to_c(args.absolute_cold_threshold_f),
    )
    write_annual_counts_csv(
        path=args.out,
        annual_counts=annual_counts,
        thresholds=thresholds,
        county_meta=county_meta,
        baseline_start_year=args.baseline_start_year,
        baseline_end_year=args.baseline_end_year,
        absolute_heat_threshold_f=args.absolute_heat_threshold_f,
        absolute_cold_threshold_f=args.absolute_cold_threshold_f,
    )

    if not args.skip_comparison:
        write_comparison_csv(
            path=args.comparison_out,
            annual_counts=annual_counts,
            county_meta=county_meta,
            old_climate_data=args.old_climate_data,
            absolute_heat_threshold_f=args.absolute_heat_threshold_f,
            absolute_cold_threshold_f=args.absolute_cold_threshold_f,
        )


if __name__ == "__main__":
    main()
