#!/usr/bin/env python3
"""
Summarize downloaded NSRDB county polygon GHI archives.

Each downloaded polygon archive contains one CSV per NSRDB site that intersects
the county polygon or tile. This script computes area-weighted county-level
average daily GHI values across those site CSVs, combining multiple tile
archives back into one county summary when present.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import statistics
import zipfile
from pathlib import Path


DEFAULT_ARCHIVE_DIR = Path("data/nrel/polygon_archives")
DEFAULT_COUNTIES_GEOJSON = Path("data/geojson-counties-fips.json")
DEFAULT_REQUESTS_CSV = Path("data/nrel/county_polygon_ghi_request_manifest.csv")
DEFAULT_REPRESENTATIVE_POINT_CSV = Path("data/nrel/county_representative_point_ghi_summary.csv")
DEFAULT_OUTPUT_CSV = Path("data/nrel/county_polygon_ghi_summary.csv")
DEFAULT_CELL_SIZE_M = 4000.0
DEFAULT_ARCHIVE_GLOB = "*_ghi.zip"
EARTH_RADIUS_M = 6_371_008.8

SUMMARY_FIELDS = [
    "county_fips",
    "county_name",
    "state_abbr",
    "archive_zip",
    "polygon_sites",
    "request_site_count",
    "ghi_rows_per_site_min",
    "ghi_rows_per_site_max",
    "avgSolarGhiKwhM2Day",
    "areaWeightedAvgSolarGhiKwhM2Day",
    "areaWeightedSites",
    "weightedCellAreaKm2",
    "countyAreaKm2",
    "site_avg_min",
    "site_avg_max",
    "site_avg_stddev",
    "representativePointAvgSolarGhiKwhM2Day",
    "areaWeightedMinusRepresentativePoint",
    "areaWeightedPctDiffFromRepresentativePoint",
]


def read_lookup(csv_path: Path, key_field: str) -> dict[str, dict[str, str]]:
    """Read a CSV into a dictionary keyed by one field."""
    if not csv_path.exists():
        return {}

    with csv_path.open(newline="", encoding="utf-8-sig") as handle:
        return {
            row[key_field]: row
            for row in csv.DictReader(handle)
            if row.get(key_field)
        }


def read_grouped_lookup(csv_path: Path, key_field: str) -> dict[str, list[dict[str, str]]]:
    """Read a CSV into lists of rows keyed by one field."""
    rows_by_key: dict[str, list[dict[str, str]]] = {}
    if not csv_path.exists():
        return rows_by_key

    with csv_path.open(newline="", encoding="utf-8-sig") as handle:
        for row in csv.DictReader(handle):
            key = row.get(key_field)
            if key:
                rows_by_key.setdefault(key, []).append(row)
    return rows_by_key


def county_fips_from_archive(path: Path) -> str:
    """Extract the county FIPS prefix from an archive filename."""
    return path.name.split("_", 1)[0]


def archive_groups_by_county(archives: list[Path]) -> dict[str, list[Path]]:
    """Group archive ZIP paths by county FIPS."""
    grouped: dict[str, list[Path]] = {}
    for archive in archives:
        grouped.setdefault(county_fips_from_archive(archive), []).append(archive)
    return grouped


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 load_county_geometries(counties_geojson: Path) -> dict[str, list[list[list[tuple[float, float]]]]]:
    """Load county polygon rings from GeoJSON without third-party GIS dependencies."""
    with counties_geojson.open(encoding="utf-8") as handle:
        geojson = json.load(handle)

    county_geometries: dict[str, list[list[list[tuple[float, float]]]]] = {}
    for feature in geojson.get("features", []):
        properties = feature.get("properties", {})
        county_fips = normalize_fips(properties.get("id"), 5)
        if not county_fips:
            county_fips = normalize_fips(properties.get("STATE"), 2) + normalize_fips(properties.get("COUNTY"), 3)
        geometry = feature.get("geometry") or {}
        polygons = geometry_polygons(geometry)
        if county_fips and polygons:
            county_geometries[county_fips] = polygons

    return county_geometries


def geometry_polygons(geometry: dict[str, object]) -> list[list[list[tuple[float, float]]]]:
    """Return GeoJSON Polygon/MultiPolygon coordinates as polygon rings."""
    geometry_type = geometry.get("type")
    coordinates = geometry.get("coordinates")
    if not isinstance(coordinates, list):
        return []

    if geometry_type == "Polygon":
        return [coordinates_to_polygon(coordinates)]
    if geometry_type == "MultiPolygon":
        return [coordinates_to_polygon(polygon) for polygon in coordinates]
    return []


def coordinates_to_polygon(coordinates: list[object]) -> list[list[tuple[float, float]]]:
    """Convert raw GeoJSON polygon coordinates to typed rings."""
    polygon: list[list[tuple[float, float]]] = []
    for ring in coordinates:
        polygon.append([(float(point[0]), float(point[1])) for point in ring])
    return polygon


def extract_site_lon_lat(csv_text: str) -> tuple[float, float]:
    """Extract the NSRDB site longitude and latitude from CSV metadata rows."""
    rows = list(csv.reader(csv_text.splitlines()))
    if len(rows) < 2:
        raise ValueError("Could not read NSRDB metadata rows.")

    metadata = {key.strip().lower(): value.strip() for key, value in zip(rows[0], rows[1])}
    try:
        lon = float(metadata["longitude"])
        lat = float(metadata["latitude"])
    except KeyError as error:
        raise ValueError("Could not find Longitude/Latitude metadata in NSRDB CSV.") from error
    except ValueError as error:
        raise ValueError("NSRDB Longitude/Latitude metadata is not numeric.") from error

    return lon, lat


def extract_ghi_values(csv_text: str) -> list[float]:
    """Extract numeric GHI values from an NSRDB CSV response."""
    rows = list(csv.reader(csv_text.splitlines()))
    header_index = next(
        (
            index
            for index, row in enumerate(rows)
            if {"Year", "Month", "Day", "Hour", "Minute", "GHI"}.issubset(set(row))
        ),
        None,
    )

    if header_index is None:
        raise ValueError("Could not find the NSRDB data header row with a GHI column.")

    ghi_index = rows[header_index].index("GHI")
    ghi_values: list[float] = []
    for row in rows[header_index + 1 :]:
        if len(row) <= ghi_index or not row[ghi_index].strip():
            continue
        ghi_values.append(float(row[ghi_index]))

    if not ghi_values:
        raise ValueError("No GHI values were found in the NSRDB response.")

    return ghi_values


def site_average_daily_ghi(csv_text: str) -> tuple[float, int]:
    """Return average daily GHI in kWh/m2/day and number of GHI rows."""
    ghi_values = extract_ghi_values(csv_text)
    return sum(ghi_values) / 1000 / 365, len(ghi_values)


def average_geometry_latitude(county_geometry: list[list[list[tuple[float, float]]]]) -> float:
    """Return a representative latitude for local meter projection."""
    latitudes = [
        lat
        for polygon in county_geometry
        for ring in polygon[:1]
        for _lon, lat in ring
    ]
    if not latitudes:
        return 0.0
    return statistics.fmean(latitudes)


def project_lon_lat(lon: float, lat: float, reference_lat: float) -> tuple[float, float]:
    """Project lon/lat to local meters with an equirectangular approximation."""
    reference_lat_radians = math.radians(reference_lat)
    x = EARTH_RADIUS_M * math.radians(lon) * math.cos(reference_lat_radians)
    y = EARTH_RADIUS_M * math.radians(lat)
    return x, y


def project_county_geometry(
    county_geometry: list[list[list[tuple[float, float]]]],
    reference_lat: float,
) -> list[list[list[tuple[float, float]]]]:
    """Project county polygon rings to local meter coordinates."""
    return [
        [
            [project_lon_lat(lon, lat, reference_lat) for lon, lat in ring]
            for ring in polygon
        ]
        for polygon in county_geometry
    ]


def polygon_area(points: list[tuple[float, float]]) -> float:
    """Return absolute polygon area using the shoelace formula."""
    if len(points) < 3:
        return 0.0

    area = 0.0
    for index, (x1, y1) in enumerate(points):
        x2, y2 = points[(index + 1) % len(points)]
        area += x1 * y2 - x2 * y1
    return abs(area) / 2


def remove_closing_point(points: list[tuple[float, float]]) -> list[tuple[float, float]]:
    """Remove duplicate final point before polygon clipping."""
    if len(points) > 1 and points[0] == points[-1]:
        return points[:-1]
    return points


def clip_ring_to_rectangle(
    ring: list[tuple[float, float]],
    min_x: float,
    min_y: float,
    max_x: float,
    max_y: float,
) -> list[tuple[float, float]]:
    """Clip a polygon ring to an axis-aligned rectangle."""
    points = remove_closing_point(ring)

    def clip_edge(
        input_points: list[tuple[float, float]],
        inside: object,
        intersect: object,
    ) -> list[tuple[float, float]]:
        if not input_points:
            return []

        output_points: list[tuple[float, float]] = []
        previous = input_points[-1]
        previous_inside = inside(previous)
        for current in input_points:
            current_inside = inside(current)
            if current_inside:
                if not previous_inside:
                    output_points.append(intersect(previous, current))
                output_points.append(current)
            elif previous_inside:
                output_points.append(intersect(previous, current))
            previous = current
            previous_inside = current_inside
        return output_points

    def vertical_intersection(x_value: float, start: tuple[float, float], end: tuple[float, float]) -> tuple[float, float]:
        x1, y1 = start
        x2, y2 = end
        if x2 == x1:
            return x_value, y1
        ratio = (x_value - x1) / (x2 - x1)
        return x_value, y1 + ratio * (y2 - y1)

    def horizontal_intersection(y_value: float, start: tuple[float, float], end: tuple[float, float]) -> tuple[float, float]:
        x1, y1 = start
        x2, y2 = end
        if y2 == y1:
            return x1, y_value
        ratio = (y_value - y1) / (y2 - y1)
        return x1 + ratio * (x2 - x1), y_value

    points = clip_edge(points, lambda point: point[0] >= min_x, lambda start, end: vertical_intersection(min_x, start, end))
    points = clip_edge(points, lambda point: point[0] <= max_x, lambda start, end: vertical_intersection(max_x, start, end))
    points = clip_edge(points, lambda point: point[1] >= min_y, lambda start, end: horizontal_intersection(min_y, start, end))
    points = clip_edge(points, lambda point: point[1] <= max_y, lambda start, end: horizontal_intersection(max_y, start, end))
    return points


def projected_polygon_area(polygon: list[list[tuple[float, float]]]) -> float:
    """Return projected polygon area, subtracting interior rings."""
    if not polygon:
        return 0.0

    area = polygon_area(polygon[0])
    for hole in polygon[1:]:
        area -= polygon_area(hole)
    return max(0.0, area)


def clipped_projected_polygon_area(
    polygon: list[list[tuple[float, float]]],
    min_x: float,
    min_y: float,
    max_x: float,
    max_y: float,
) -> float:
    """Return projected polygon area inside a rectangle, subtracting holes."""
    if not polygon:
        return 0.0

    clipped_exterior = clip_ring_to_rectangle(polygon[0], min_x, min_y, max_x, max_y)
    area = polygon_area(clipped_exterior)
    for hole in polygon[1:]:
        clipped_hole = clip_ring_to_rectangle(hole, min_x, min_y, max_x, max_y)
        area -= polygon_area(clipped_hole)
    return max(0.0, area)


def area_weighted_average(
    site_averages: list[float],
    site_lon_lats: list[tuple[float, float]],
    county_geometry: list[list[list[tuple[float, float]]]] | None,
    cell_size_m: float,
) -> dict[str, float | int]:
    """Compute a county average weighted by estimated grid-cell overlap area."""
    if county_geometry is None:
        return {
            "area_weighted_avg": math.nan,
            "area_weighted_sites": 0,
            "weighted_cell_area_km2": math.nan,
            "county_area_km2": math.nan,
        }

    reference_lat = average_geometry_latitude(county_geometry)
    projected_county = project_county_geometry(county_geometry, reference_lat)
    half_cell = cell_size_m / 2
    weighted_sum = 0.0
    total_weight = 0.0
    weighted_sites = 0
    county_area = sum(projected_polygon_area(polygon) for polygon in projected_county)

    for site_average, (lon, lat) in zip(site_averages, site_lon_lats):
        x, y = project_lon_lat(lon, lat, reference_lat)
        min_x = x - half_cell
        min_y = y - half_cell
        max_x = x + half_cell
        max_y = y + half_cell
        overlap_area = sum(
            clipped_projected_polygon_area(polygon, min_x, min_y, max_x, max_y)
            for polygon in projected_county
        )
        if overlap_area <= 0:
            continue
        weighted_sum += site_average * overlap_area
        total_weight += overlap_area
        weighted_sites += 1

    area_weighted_avg = weighted_sum / total_weight if total_weight else math.nan
    return {
        "area_weighted_avg": area_weighted_avg,
        "area_weighted_sites": weighted_sites,
        "weighted_cell_area_km2": total_weight / 1_000_000,
        "county_area_km2": county_area / 1_000_000,
    }


def summarize_archives(
    paths: list[Path],
    county_geometry: list[list[list[tuple[float, float]]]] | None,
    cell_size_m: float,
) -> dict[str, float | int | str]:
    """Summarize all NSRDB site CSVs in one or more ZIP archives."""
    site_averages: list[float] = []
    site_lon_lats: list[tuple[float, float]] = []
    row_counts: list[int] = []

    for path in paths:
        with zipfile.ZipFile(path) as archive:
            csv_names = sorted(name for name in archive.namelist() if name.lower().endswith(".csv"))
            for csv_name in csv_names:
                csv_text = archive.read(csv_name).decode("utf-8-sig")
                site_average, row_count = site_average_daily_ghi(csv_text)
                site_lon_lat = extract_site_lon_lat(csv_text)
                site_averages.append(site_average)
                site_lon_lats.append(site_lon_lat)
                row_counts.append(row_count)

    if not site_averages:
        raise ValueError(f"No CSV files found in {', '.join(str(path) for path in paths)}")

    weighted = area_weighted_average(site_averages, site_lon_lats, county_geometry, cell_size_m)

    return {
        "polygon_sites": len(site_averages),
        "ghi_rows_per_site_min": min(row_counts),
        "ghi_rows_per_site_max": max(row_counts),
        "avgSolarGhiKwhM2Day": weighted["area_weighted_avg"],
        "area_weighted_avg": weighted["area_weighted_avg"],
        "area_weighted_sites": weighted["area_weighted_sites"],
        "weighted_cell_area_km2": weighted["weighted_cell_area_km2"],
        "county_area_km2": weighted["county_area_km2"],
        "site_avg_min": min(site_averages),
        "site_avg_max": max(site_averages),
        "site_avg_stddev": statistics.pstdev(site_averages) if len(site_averages) > 1 else 0.0,
    }


def summarize_archive(
    path: Path,
    county_geometry: list[list[list[tuple[float, float]]]] | None,
    cell_size_m: float,
) -> dict[str, float | int | str]:
    """Summarize all NSRDB site CSVs in one ZIP archive."""
    return summarize_archives([path], county_geometry, cell_size_m)


def format_float(value: float | None, places: int = 3) -> str:
    """Format an optional float for CSV output."""
    if value is None or math.isnan(value):
        return ""
    return f"{value:.{places}f}"


def build_summary_row(
    archive_paths: list[Path],
    polygon_summary: dict[str, float | int | str],
    request_rows: list[dict[str, str]],
    representative_point_row: dict[str, str],
) -> dict[str, str]:
    """Build one output row from polygon and optional representative-point summaries."""
    county_fips = county_fips_from_archive(archive_paths[0])
    request_row = request_rows[0] if request_rows else {}
    request_site_counts = [
        int(row["site_count"])
        for row in request_rows
        if str(row.get("site_count", "")).strip().isdigit()
    ]
    archive_zip = ";".join(str(path) for path in archive_paths)
    final_avg = float(polygon_summary["avgSolarGhiKwhM2Day"])
    area_weighted_avg = float(polygon_summary["area_weighted_avg"])
    representative_point_avg = (
        float(representative_point_row["avgSolarGhiKwhM2Day"])
        if representative_point_row.get("avgSolarGhiKwhM2Day")
        else None
    )
    weighted_minus_representative_point = (
        area_weighted_avg - representative_point_avg if representative_point_avg is not None else None
    )
    weighted_pct_diff_representative_point = (
        weighted_minus_representative_point / representative_point_avg * 100
        if representative_point_avg not in (None, 0)
        else None
    )

    return {
        "county_fips": county_fips,
        "county_name": request_row.get("county_name") or representative_point_row.get("county_name", ""),
        "state_abbr": request_row.get("state_abbr") or representative_point_row.get("state_abbr", ""),
        "archive_zip": archive_zip,
        "polygon_sites": str(polygon_summary["polygon_sites"]),
        "request_site_count": str(sum(request_site_counts)) if request_site_counts else request_row.get("site_count", ""),
        "ghi_rows_per_site_min": str(polygon_summary["ghi_rows_per_site_min"]),
        "ghi_rows_per_site_max": str(polygon_summary["ghi_rows_per_site_max"]),
        "avgSolarGhiKwhM2Day": format_float(final_avg),
        "areaWeightedAvgSolarGhiKwhM2Day": format_float(area_weighted_avg),
        "areaWeightedSites": str(polygon_summary["area_weighted_sites"]),
        "weightedCellAreaKm2": format_float(float(polygon_summary["weighted_cell_area_km2"]), 1),
        "countyAreaKm2": format_float(float(polygon_summary["county_area_km2"]), 1),
        "site_avg_min": format_float(float(polygon_summary["site_avg_min"])),
        "site_avg_max": format_float(float(polygon_summary["site_avg_max"])),
        "site_avg_stddev": format_float(float(polygon_summary["site_avg_stddev"])),
        "representativePointAvgSolarGhiKwhM2Day": format_float(representative_point_avg),
        "areaWeightedMinusRepresentativePoint": format_float(weighted_minus_representative_point),
        "areaWeightedPctDiffFromRepresentativePoint": format_float(weighted_pct_diff_representative_point),
    }


def archive_signature(paths: list[Path]) -> set[str]:
    """Return archive filenames for comparing existing summary rows."""
    return {path.name for path in paths}


def existing_archive_signature(row: dict[str, str]) -> set[str]:
    """Return archive filenames recorded in an existing summary row."""
    return {
        Path(path_text).name
        for path_text in row.get("archive_zip", "").split(";")
        if path_text.strip()
    }


def write_summary(rows: list[dict[str, str]], output_csv: Path) -> None:
    """Write polygon GHI summary rows to disk."""
    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=SUMMARY_FIELDS)
        writer.writeheader()
        writer.writerows(rows)


def run(args: argparse.Namespace) -> None:
    """Summarize downloaded polygon archives."""
    county_geometries = load_county_geometries(args.counties_geojson)
    request_rows = read_grouped_lookup(args.requests_csv, "county_fips")
    representative_point_rows = read_lookup(args.representative_point_csv, "county_fips")
    existing_rows = (
        read_lookup(args.output_csv, "county_fips")
        if args.reuse_existing_output
        else {}
    )
    archives = sorted(args.archive_dir.glob(args.archive_glob))

    if not archives:
        print(f"No archives matched {args.archive_dir / args.archive_glob}")
        return

    rows: list[dict[str, str]] = []
    archive_groups = archive_groups_by_county(archives)
    total = len(archive_groups)
    reused_count = 0
    for index, county_fips in enumerate(sorted(archive_groups), start=1):
        archive_paths = sorted(archive_groups[county_fips])
        existing_row = existing_rows.get(county_fips)
        if existing_row and existing_archive_signature(existing_row) == archive_signature(archive_paths):
            rows.append({field: existing_row.get(field, "") for field in SUMMARY_FIELDS})
            reused_count += 1
            continue

        try:
            polygon_summary = summarize_archives(
                archive_paths,
                county_geometries.get(county_fips),
                args.cell_size_m,
            )
            row = build_summary_row(
                archive_paths,
                polygon_summary,
                request_rows.get(county_fips, []),
                representative_point_rows.get(county_fips, {}),
            )
            rows.append(row)
        except (OSError, ValueError, zipfile.BadZipFile) as error:
            print(f"[{index}/{total}] Failed {county_fips}: {error}")
            continue

        tile_note = f", archives={len(archive_paths)}" if len(archive_paths) > 1 else ""
        print(
            f"[{index}/{total}] {county_fips}: "
            f"sites={row['polygon_sites']}, weighted={row['areaWeightedAvgSolarGhiKwhM2Day']}, "
            f"representative_point={row['representativePointAvgSolarGhiKwhM2Day'] or 'n/a'}"
            f"{tile_note}"
        )

    write_summary(rows, args.output_csv)
    if reused_count:
        print(f"Reused {reused_count} existing rows whose archive sets still match.")
    print(f"Wrote {len(rows)} rows to {args.output_csv}")


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--archive-dir", type=Path, default=DEFAULT_ARCHIVE_DIR)
    parser.add_argument("--archive-glob", default=DEFAULT_ARCHIVE_GLOB)
    parser.add_argument("--counties-geojson", type=Path, default=DEFAULT_COUNTIES_GEOJSON)
    parser.add_argument("--requests-csv", type=Path, default=DEFAULT_REQUESTS_CSV)
    parser.add_argument(
        "--representative-point-csv",
        type=Path,
        default=DEFAULT_REPRESENTATIVE_POINT_CSV,
        help="County representative-point GHI summary CSV.",
    )
    parser.add_argument("--output-csv", type=Path, default=DEFAULT_OUTPUT_CSV)
    parser.add_argument(
        "--cell-size-m",
        type=float,
        default=DEFAULT_CELL_SIZE_M,
        help="Estimated NSRDB grid-cell side length in meters. GOES TMY is 4 km.",
    )
    parser.add_argument(
        "--reuse-existing-output",
        action="store_true",
        help=(
            "Reuse rows already present in --output-csv when the recorded archive filenames "
            "match the current archive directory, and summarize only missing or changed counties."
        ),
    )
    args = parser.parse_args()
    if args.cell_size_m <= 0:
        parser.error("--cell-size-m must be greater than 0.")
    return args


def main() -> None:
    """Run the polygon GHI summarizer."""
    run(parse_args())


if __name__ == "__main__":
    main()
