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

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

Primary metric:

cloudinessIndexPct = 1 - mean(clamped(GHI / Clearsky GHI, 0, 1))

Despite the legacy "Pct" field name, the primary index is stored on a 0-1 scale.
Cloud Type bucket fields are stored as percentages.
"""

from __future__ import annotations

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

from summarize_nsrdb_county_polygon_archives import (
    area_weighted_average,
    archive_groups_by_county,
    archive_signature,
    county_fips_from_archive,
    existing_archive_signature,
    extract_site_lon_lat,
    format_float,
    load_county_geometries,
    read_grouped_lookup,
    read_lookup,
)


DEFAULT_ARCHIVE_DIR = Path("data/nrel/polygon_cloud_archives")
DEFAULT_COUNTIES_GEOJSON = Path("data/geojson-counties-fips.json")
DEFAULT_REQUESTS_CSV = Path("data/nrel/county_polygon_cloud_request_manifest.csv")
DEFAULT_REPRESENTATIVE_POINT_CSV = Path("data/nrel/county_representative_point_cloud_summary.csv")
DEFAULT_OUTPUT_CSV = Path("data/nrel/county_polygon_cloud_summary.csv")
DEFAULT_CELL_SIZE_M = 4000.0
DEFAULT_ARCHIVE_GLOB = "*.zip"
DEFAULT_MIN_CLEARSKY_GHI = 50.0

SUMMARY_FIELDS = [
    "county_fips",
    "county_name",
    "state_abbr",
    "archive_zip",
    "polygon_sites",
    "request_site_count",
    "rows_per_site_min",
    "rows_per_site_max",
    "daylight_rows_per_site_min",
    "daylight_rows_per_site_max",
    "cloudinessIndexPct",
    "areaWeightedCloudinessIndexPct",
    "areaWeightedAvgObservedToClearskyRatio",
    "areaWeightedSites",
    "weightedCellAreaKm2",
    "countyAreaKm2",
    "site_cloudiness_min",
    "site_cloudiness_max",
    "site_cloudiness_stddev",
    "clearOrProbablyClearPct",
    "cloudyOrObscuredPct",
    "fogPct",
    "waterCloudPct",
    "iceCloudPct",
    "cirrusPct",
    "unknownCloudTypePct",
    "representativePointCloudinessIndexPct",
    "areaWeightedMinusRepresentativePoint",
    "areaWeightedPctDiffFromRepresentativePoint",
]


def normalized_header(value: str) -> str:
    """Normalize a CSV header to make NSRDB spelling variations easier to match."""
    return re.sub(r"[^a-z0-9]+", "", value.lower())


def find_column(header: list[str], candidates: set[str]) -> int | None:
    """Return the first header index whose normalized name matches a candidate."""
    normalized_candidates = {normalized_header(candidate) for candidate in candidates}
    for index, value in enumerate(header):
        if normalized_header(value) in normalized_candidates:
            return index
    return None


def find_data_header(rows: list[list[str]]) -> int | None:
    """Return the NSRDB hourly data header row index."""
    required = {"Year", "Month", "Day", "Hour", "Minute"}
    for index, row in enumerate(rows):
        if required.issubset(set(row)):
            return index
    return None


def parse_optional_float(row: list[str], index: int | None) -> float | None:
    """Parse an optional float field from a CSV row."""
    if index is None or len(row) <= index:
        return None
    value = row[index].strip()
    if not value:
        return None
    try:
        parsed = float(value)
    except ValueError:
        return None
    return parsed if math.isfinite(parsed) else None


def parse_cloud_type(row: list[str], index: int | None) -> str:
    """Parse an optional Cloud Type code as a stable string."""
    if index is None or len(row) <= index:
        return ""
    value = row[index].strip()
    if not value:
        return ""
    try:
        return str(int(float(value)))
    except ValueError:
        return value


def pct(count: int, total: int) -> float:
    """Return a percentage or NaN when there is no denominator."""
    return count / total * 100 if total else math.nan


def site_cloud_metrics(csv_text: str, min_clearsky_ghi: float) -> dict[str, float | int]:
    """Return site-level cloudiness metrics from one NSRDB CSV response."""
    rows = list(csv.reader(csv_text.splitlines()))
    header_index = find_data_header(rows)
    if header_index is None:
        raise ValueError("Could not find the NSRDB data header row.")

    header = rows[header_index]
    ghi_index = find_column(header, {"GHI", "ghi"})
    clearsky_ghi_index = find_column(
        header,
        {
            "Clearsky GHI",
            "Clear Sky GHI",
            "Clear-sky GHI",
            "clearsky_ghi",
            "clear_sky_ghi",
        },
    )
    cloud_type_index = find_column(header, {"Cloud Type", "cloud_type", "cloudtype"})
    if ghi_index is None:
        raise ValueError("Could not find a GHI column in the NSRDB response.")
    if clearsky_ghi_index is None:
        raise ValueError("Could not find a Clearsky GHI column in the NSRDB response.")
    if cloud_type_index is None:
        raise ValueError("Could not find a Cloud Type column in the NSRDB response.")

    all_rows = 0
    daylight_rows = 0
    ratio_sum = 0.0
    daylight_cloud_type_counts: dict[str, int] = {}

    for row in rows[header_index + 1 :]:
        if not row:
            continue
        all_rows += 1
        ghi = parse_optional_float(row, ghi_index)
        clearsky_ghi = parse_optional_float(row, clearsky_ghi_index)
        cloud_type = parse_cloud_type(row, cloud_type_index)
        if ghi is None or clearsky_ghi is None or clearsky_ghi < min_clearsky_ghi:
            continue

        ratio_sum += min(1.0, max(0.0, ghi / clearsky_ghi))
        daylight_rows += 1
        if cloud_type:
            daylight_cloud_type_counts[cloud_type] = daylight_cloud_type_counts.get(cloud_type, 0) + 1

    if daylight_rows == 0:
        raise ValueError("No daylight rows with valid GHI and Clearsky GHI were found.")

    avg_ratio = ratio_sum / daylight_rows
    cloudiness = 1 - avg_ratio
    clear_or_probably_clear = daylight_cloud_type_counts.get("0", 0) + daylight_cloud_type_counts.get("1", 0)
    cloudy_or_obscured = sum(
        daylight_cloud_type_counts.get(code, 0)
        for code in ["2", "3", "4", "5", "6", "7", "8", "9", "11", "12"]
    )
    water_cloud = (
        daylight_cloud_type_counts.get("3", 0)
        + daylight_cloud_type_counts.get("4", 0)
        + daylight_cloud_type_counts.get("5", 0)
    )
    ice_cloud = (
        daylight_cloud_type_counts.get("6", 0)
        + daylight_cloud_type_counts.get("8", 0)
        + daylight_cloud_type_counts.get("9", 0)
    )
    unknown_cloud_type = daylight_cloud_type_counts.get("10", 0) + daylight_cloud_type_counts.get("-15", 0)

    return {
        "cloudiness": cloudiness,
        "avg_ratio": avg_ratio,
        "all_rows": all_rows,
        "daylight_rows": daylight_rows,
        "clear_or_probably_clear_pct": pct(clear_or_probably_clear, daylight_rows),
        "cloudy_or_obscured_pct": pct(cloudy_or_obscured, daylight_rows),
        "fog_pct": pct(daylight_cloud_type_counts.get("2", 0), daylight_rows),
        "water_cloud_pct": pct(water_cloud, daylight_rows),
        "ice_cloud_pct": pct(ice_cloud, daylight_rows),
        "cirrus_pct": pct(daylight_cloud_type_counts.get("7", 0), daylight_rows),
        "unknown_cloud_type_pct": pct(unknown_cloud_type, daylight_rows),
    }


def weighted_metric(
    site_metrics: list[dict[str, float | int]],
    metric_key: str,
    site_lon_lats: list[tuple[float, float]],
    county_geometry: list[list[list[tuple[float, float]]]] | None,
    cell_size_m: float,
) -> dict[str, float | int]:
    """Area-weight one metric across sites using estimated grid-cell overlap."""
    values = [float(metrics[metric_key]) for metrics in site_metrics]
    return area_weighted_average(values, site_lon_lats, county_geometry, cell_size_m)


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

    for path in paths:
        with zipfile.ZipFile(path) as archive:
            bad_member = archive.testzip()
            if bad_member:
                raise zipfile.BadZipFile(
                    f"CRC check failed for {bad_member} in {path.name}"
                )
            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")
                if not csv_text.strip():
                    empty_csv_members.append(f"{path.name}:{csv_name}")
                    continue
                site_metrics.append(site_cloud_metrics(csv_text, min_clearsky_ghi))
                site_lon_lats.append(extract_site_lon_lat(csv_text))

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

    cloudiness_values = [float(metrics["cloudiness"]) for metrics in site_metrics]
    row_counts = [int(metrics["all_rows"]) for metrics in site_metrics]
    daylight_row_counts = [int(metrics["daylight_rows"]) for metrics in site_metrics]

    weighted_cloudiness = weighted_metric(site_metrics, "cloudiness", site_lon_lats, county_geometry, cell_size_m)
    weighted_avg_ratio = weighted_metric(site_metrics, "avg_ratio", site_lon_lats, county_geometry, cell_size_m)

    return {
        "polygon_sites": len(site_metrics),
        "rows_per_site_min": min(row_counts),
        "rows_per_site_max": max(row_counts),
        "daylight_rows_per_site_min": min(daylight_row_counts),
        "daylight_rows_per_site_max": max(daylight_row_counts),
        "area_weighted_cloudiness": weighted_cloudiness["area_weighted_avg"],
        "area_weighted_avg_ratio": weighted_avg_ratio["area_weighted_avg"],
        "area_weighted_sites": weighted_cloudiness["area_weighted_sites"],
        "weighted_cell_area_km2": weighted_cloudiness["weighted_cell_area_km2"],
        "county_area_km2": weighted_cloudiness["county_area_km2"],
        "site_cloudiness_min": min(cloudiness_values),
        "site_cloudiness_max": max(cloudiness_values),
        "site_cloudiness_stddev": statistics.pstdev(cloudiness_values) if len(cloudiness_values) > 1 else 0.0,
        "clear_or_probably_clear_pct": weighted_metric(
            site_metrics, "clear_or_probably_clear_pct", site_lon_lats, county_geometry, cell_size_m
        )["area_weighted_avg"],
        "cloudy_or_obscured_pct": weighted_metric(
            site_metrics, "cloudy_or_obscured_pct", site_lon_lats, county_geometry, cell_size_m
        )["area_weighted_avg"],
        "fog_pct": weighted_metric(site_metrics, "fog_pct", site_lon_lats, county_geometry, cell_size_m)[
            "area_weighted_avg"
        ],
        "water_cloud_pct": weighted_metric(
            site_metrics, "water_cloud_pct", site_lon_lats, county_geometry, cell_size_m
        )["area_weighted_avg"],
        "ice_cloud_pct": weighted_metric(site_metrics, "ice_cloud_pct", site_lon_lats, county_geometry, cell_size_m)[
            "area_weighted_avg"
        ],
        "cirrus_pct": weighted_metric(site_metrics, "cirrus_pct", site_lon_lats, county_geometry, cell_size_m)[
            "area_weighted_avg"
        ],
        "unknown_cloud_type_pct": weighted_metric(
            site_metrics, "unknown_cloud_type_pct", site_lon_lats, county_geometry, cell_size_m
        )["area_weighted_avg"],
        "empty_csv_members": len(empty_csv_members),
    }


def archive_part_id(path: Path) -> str:
    """Return county or tileNNN for one downloaded archive filename."""
    parts = path.stem.split("_")
    return parts[1] if len(parts) > 1 and parts[1].startswith("tile") else "county"


def missing_archive_parts(
    archive_paths: list[Path],
    request_rows: list[dict[str, str]],
) -> list[str]:
    """Return requested county/tile archive parts that are not downloaded."""
    expected = {
        row.get("tile_id", "").strip() or "county"
        for row in request_rows
        if row.get("download_url", "").strip()
    }
    actual = {archive_part_id(path) for path in archive_paths}
    return sorted(expected - actual)


def build_summary_row(
    archive_paths: list[Path],
    polygon_summary: dict[str, float | int],
    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)
    area_weighted_cloudiness = float(polygon_summary["area_weighted_cloudiness"])
    representative_point_cloudiness = (
        float(representative_point_row["cloudinessIndexPct"])
        if representative_point_row.get("cloudinessIndexPct")
        else None
    )
    weighted_minus_representative_point = (
        area_weighted_cloudiness - representative_point_cloudiness
        if representative_point_cloudiness is not None
        else None
    )
    weighted_pct_diff_representative_point = (
        weighted_minus_representative_point / representative_point_cloudiness * 100
        if representative_point_cloudiness 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", ""),
        "rows_per_site_min": str(polygon_summary["rows_per_site_min"]),
        "rows_per_site_max": str(polygon_summary["rows_per_site_max"]),
        "daylight_rows_per_site_min": str(polygon_summary["daylight_rows_per_site_min"]),
        "daylight_rows_per_site_max": str(polygon_summary["daylight_rows_per_site_max"]),
        "cloudinessIndexPct": format_float(area_weighted_cloudiness, 4),
        "areaWeightedCloudinessIndexPct": format_float(area_weighted_cloudiness, 4),
        "areaWeightedAvgObservedToClearskyRatio": format_float(float(polygon_summary["area_weighted_avg_ratio"]), 4),
        "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_cloudiness_min": format_float(float(polygon_summary["site_cloudiness_min"]), 4),
        "site_cloudiness_max": format_float(float(polygon_summary["site_cloudiness_max"]), 4),
        "site_cloudiness_stddev": format_float(float(polygon_summary["site_cloudiness_stddev"]), 4),
        "clearOrProbablyClearPct": format_float(float(polygon_summary["clear_or_probably_clear_pct"]), 2),
        "cloudyOrObscuredPct": format_float(float(polygon_summary["cloudy_or_obscured_pct"]), 2),
        "fogPct": format_float(float(polygon_summary["fog_pct"]), 2),
        "waterCloudPct": format_float(float(polygon_summary["water_cloud_pct"]), 2),
        "iceCloudPct": format_float(float(polygon_summary["ice_cloud_pct"]), 2),
        "cirrusPct": format_float(float(polygon_summary["cirrus_pct"]), 2),
        "unknownCloudTypePct": format_float(float(polygon_summary["unknown_cloud_type_pct"]), 2),
        "representativePointCloudinessIndexPct": format_float(representative_point_cloudiness, 4),
        "areaWeightedMinusRepresentativePoint": format_float(weighted_minus_representative_point, 4),
        "areaWeightedPctDiffFromRepresentativePoint": format_float(weighted_pct_diff_representative_point),
    }


def write_summary(rows: list[dict[str, str]], output_csv: Path) -> None:
    """Write polygon cloud 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 cloud 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])
        county_request_rows = request_rows.get(county_fips, [])
        missing_parts = missing_archive_parts(archive_paths, county_request_rows)
        if missing_parts:
            print(
                f"[{index}/{total}] Skipped {county_fips}: incomplete archive set; "
                f"missing {', '.join(missing_parts)}."
            )
            continue

        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,
                args.min_clearsky_ghi,
            )
            row = build_summary_row(
                archive_paths,
                polygon_summary,
                county_request_rows,
                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 ""
        empty_note = (
            f", skipped_empty_csvs={polygon_summary['empty_csv_members']}"
            if polygon_summary["empty_csv_members"]
            else ""
        )
        print(
            f"[{index}/{total}] {county_fips}: "
            f"sites={row['polygon_sites']}, weighted={row['areaWeightedCloudinessIndexPct']}, "
            f"representative_point={row['representativePointCloudinessIndexPct'] or 'n/a'}"
            f"{tile_note}{empty_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 cloud 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(
        "--min-clearsky-ghi",
        type=float,
        default=DEFAULT_MIN_CLEARSKY_GHI,
        help="Minimum Clearsky GHI W/m2 for daylight cloudiness ratio rows.",
    )
    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.")
    if args.min_clearsky_ghi < 0:
        parser.error("--min-clearsky-ghi must be 0 or greater.")
    return args


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


if __name__ == "__main__":
    main()
