#!/usr/bin/env python3
"""
Build representative county points for NSRDB point-based solar sampling.

The output CSV is intended as input to fetch_nsrdb_representative_point_ghi.py.
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

import geopandas as gpd


DEFAULT_COUNTIES_GEOJSON = Path("data/geojson-counties-fips.json")
DEFAULT_OUTPUT_CSV = Path("data/nrel/county_representative_points.csv")

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",
}


def normalize_fips(value: object, width: int) -> str:
    """Return a zero-padded FIPS string from a mixed text or numeric value."""
    digits = "".join(character for character in str(value).strip() if character.isdigit())
    return digits.zfill(width)[-width:] if digits else ""


def build_county_points(input_geojson: Path, include_puerto_rico: bool) -> list[dict[str, str]]:
    """Load county polygons and return one interior representative point per county."""
    counties = gpd.read_file(input_geojson)
    rows: list[dict[str, str]] = []

    for _, county in counties.iterrows():
        state_fips = normalize_fips(county.get("STATE"), 2)
        county_code = normalize_fips(county.get("COUNTY"), 3)
        county_fips = normalize_fips(county.get("id") or f"{state_fips}{county_code}", 5)

        if not county_fips or (state_fips == "72" and not include_puerto_rico):
            continue

        point = county.geometry.representative_point()
        county_name = str(county.get("NAME") or f"County {county_fips}").strip()

        rows.append(
            {
                "county_fips": county_fips,
                "county_name": county_name,
                "state_fips": state_fips,
                "state_abbr": STATE_FIPS_TO_ABBR.get(state_fips, state_fips),
                "lat": f"{point.y:.6f}",
                "lon": f"{point.x:.6f}",
            }
        )

    return sorted(rows, key=lambda row: row["county_fips"])


def write_points_csv(rows: list[dict[str, str]], output_csv: Path) -> None:
    """Write county representative points to CSV."""
    output_csv.parent.mkdir(parents=True, exist_ok=True)
    with output_csv.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=["county_fips", "county_name", "state_fips", "state_abbr", "lat", "lon"],
        )
        writer.writeheader()
        writer.writerows(rows)


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input-geojson", type=Path, default=DEFAULT_COUNTIES_GEOJSON)
    parser.add_argument("--output-csv", type=Path, default=DEFAULT_OUTPUT_CSV)
    parser.add_argument(
        "--include-puerto-rico",
        action="store_true",
        help="Include Puerto Rico counties. The current web app excludes Puerto Rico polygons.",
    )
    return parser.parse_args()


def main() -> None:
    """Build and write the county representative-point CSV."""
    args = parse_args()
    rows = build_county_points(args.input_geojson, args.include_puerto_rico)
    write_points_csv(rows, args.output_csv)
    print(f"Wrote {len(rows)} county representative points to {args.output_csv}")


if __name__ == "__main__":
    main()