#!/usr/bin/env python3
"""Reproduce the 2024 airport ambient screen. Python 3.9+, standard library only."""
import argparse
from collections import Counter, defaultdict
import csv
from datetime import datetime, timedelta, timezone
import hashlib
import io
import json
from pathlib import Path
from urllib.request import urlopen
from zoneinfo import ZoneInfo

SOURCE = Path(__file__).resolve()
IN_REPO = SOURCE.parent.name == "scripts"
APP = SOURCE.parents[1] if IN_REPO else Path.cwd()
ROOT = APP.parents[1] if IN_REPO else Path.cwd()
OUT = APP / "public/resources/downloads" if IN_REPO else Path.cwd() / "derived"
TZ = ZoneInfo("America/New_York")
YEAR = 2024
STATIONS = {
    "72502014734": ("Newark", "NEWARK LIBERTY INTERNATIONAL AIRPORT, NJ US", 40.68275, -74.16927),
    "72503014732": ("LaGuardia", "LAGUARDIA AIRPORT, NY US", 40.77945, -73.88027),
}
SCENARIOS = {"10C": 10, "15C": 15}
ACCEPT = {"1", "5"}
AA_FIELDS = tuple(f"AA{i}" for i in range(1, 5))
WEATHER_FIELDS = tuple(f"{prefix}{i}" for prefix in ("AU", "AW") for i in range(1, 10))
PRECIP_FIELDS = AA_FIELDS + WEATHER_FIELDS
AW_PRECIP = ({11, 21, 22, 23, 24, 25} | set(range(40, 49))
             | set(range(50, 59)) | set(range(60, 69)) | set(range(70, 79))
             | set(range(80, 88)) | {89, 92, 93, 95, 96})
AW_VALID = AW_PRECIP | {0, 1, 2, 3, 4, 5, 7, 10, 12, 18, 20, 26, 27, 28,
                       29, 30, 31, 32, 33, 34, 35, 90, 91, 94, 99}


def accumulation(value):
    """Return (hours, state). Condition 9 is not a missing numeric depth."""
    parts = value.split(",")
    if len(parts) != 4:
        return None, "unknown"
    period, depth, condition, quality = parts
    if not (len(period) == 2 and period.isdigit() and 1 <= int(period) <= 98):
        return None, "unknown"
    hours = int(period)
    if (len(depth) != 4 or not depth.isdigit() or int(depth) == 9999
            or quality not in ACCEPT or condition not in {"2", "9"}):
        return hours, "unknown"
    return hours, "precipitation" if condition == "2" or int(depth) > 0 else "zero"


def present_weather(row):
    """Positive evidence veto only; absent optional weather never proves dry."""
    states = []
    for key in WEATHER_FIELDS:
        value = row.get(key, "")
        if not value:
            continue
        parts = value.split(",")
        if key.startswith("AU"):
            valid = (len(parts) == 7 and parts[-1] in ACCEPT
                     and parts[0] in "01234" and len(parts[0]) == 1
                     and parts[1] in "012345678" and len(parts[1]) == 1
                     and parts[2] in {f"{n:02}" for n in range(10)}
                     and all(len(p) == 1 and p in "012345678" for p in parts[3:5])
                     and parts[5] in {"1", "2", "3"})
            states.append("precipitation" if valid and parts[2] != "00"
                          else "neutral" if valid else "unknown")
        else:
            valid = (len(parts) == 2 and len(parts[0]) == 2 and parts[0].isdigit()
                     and int(parts[0]) in AW_VALID and parts[1] in ACCEPT)
            states.append("precipitation" if valid and int(parts[0]) in AW_PRECIP
                          else "unknown" if not valid or int(parts[0]) in {26, 90}
                          else "neutral")
    return ("precipitation" if "precipitation" in states else
            "unknown" if "unknown" in states else "no_veto")


def precipitation_screen(bins, hours):
    """Retrospective one-hour lookback screen, independent of ambient QC."""
    # Index all routine-report accumulation intervals by overlapping UTC bin.
    # Sets deduplicate overlapping/repeated reports; amounts are NEVER summed.
    evidence = defaultdict(set)
    for reports in bins.values():
        for row in reports:
            end = datetime.fromisoformat(row["DATE"]).replace(tzinfo=timezone.utc)
            for key in AA_FIELDS:
                period, state = accumulation(row.get(key, ""))
                if period is None or state == "zero":
                    continue
                begin = end - timedelta(hours=period)
                t = begin.replace(minute=0, second=0, microsecond=0)
                while t < end:
                    evidence[t].add((begin, end, state))
                    t += timedelta(hours=1)
    result = {}
    for t in hours:
        reports = bins.get(t, [])
        if not reports:
            result[t] = "unknown"
            continue
        latest = max(r["DATE"] for r in reports)
        candidates = [r for r in reports if r["DATE"] == latest]
        if len({tuple(r.get(k, "") for k in PRECIP_FIELDS) for r in candidates}) != 1:
            result[t] = "unknown"
            continue
        row = candidates[0]
        end = datetime.fromisoformat(row["DATE"]).replace(tzinfo=timezone.utc)
        begin = end - timedelta(hours=1)
        overlaps = set()
        for bucket in (begin.replace(minute=0, second=0, microsecond=0), t):
            overlaps.update(state for a, b, state in evidence[bucket] if a < end and b > begin)
        fields = [accumulation(row[k]) for k in AA_FIELDS if row.get(k)]
        weather = present_weather(row)
        if "precipitation" in overlaps or weather == "precipitation":
            result[t] = "precipitation_evidence"
        elif ("unknown" in overlaps or weather == "unknown"
              or any(state == "unknown" for _, state in fields)):
            result[t] = "unknown"
        elif (1, "zero") in fields:
            result[t] = "reported_zero"
        else:
            result[t] = "unknown"
    return result


def measurement(value, wind=False):
    """Use tenths of C or m/s; reject missing, malformed and non-1/5 flags."""
    parts = value.split(",")
    try:
        raw, flag = (parts[3], parts[4]) if wind else (parts[0], parts[1])
        if int(raw) == 9999 or flag not in ACCEPT:
            return None
        return int(raw) / 10
    except (ValueError, IndexError):
        return None


def values(row):
    result = (measurement(row["TMP"]), measurement(row["DEW"]), measurement(row["WND"], True))
    if any(v is None for v in result):
        return None
    t, d, w = result
    return result if -93.2 <= t <= 61.8 and -98.2 <= d <= 36.8 and d <= t and 0 <= w <= 90 else None


def qualifies(result, minimum):
    return result is not None and minimum <= result[0] <= 35 and result[0] - result[1] >= 3 and result[2] <= 4.5


def long_runs(slots):
    """Return maximal runs >=4 hourly bins; never bridge a day, gap or failed bin."""
    runs, run = [], []
    for stamp, passed in slots:
        if run and (not passed or stamp - run[-1] != timedelta(hours=1)
                    or stamp.astimezone(TZ).date() != run[-1].astimezone(TZ).date()):
            if len(run) >= 4:
                runs.append(run)
            run = []
        if passed:
            run.append(stamp)
    if len(run) >= 4:
        runs.append(run)
    return runs


def choose(rows):
    """Latest routine report; conflicting selected-time measurements are unusable."""
    if not rows:
        return None
    latest = max(r["DATE"] for r in rows)
    candidates = [r for r in rows if r["DATE"] == latest]
    if len({(r["TMP"], r["DEW"], r["WND"]) for r in candidates}) != 1:
        return None
    return candidates[0]


def analyze(station, blob):
    label, name, lat, lon = STATIONS[station]
    text = blob.decode("utf-8-sig")
    reader = csv.DictReader(io.StringIO(text))
    required = {"STATION", "DATE", "NAME", "LATITUDE", "LONGITUDE", "REPORT_TYPE", "TMP", "DEW", "WND"}
    if not required.issubset(reader.fieldnames or []):
        raise ValueError("Not an ISD CSV (possibly an HTML service migration page)")
    rows = list(reader)
    if not rows:
        raise ValueError("Empty station file")
    for row in rows:
        if row["STATION"] != station or row["NAME"] != name:
            raise ValueError(f"Unexpected station identity: {station}")
        if abs(float(row["LATITUDE"]) - lat) > .02 or abs(float(row["LONGITUDE"]) - lon) > .02:
            raise ValueError(f"Unexpected station coordinates: {station}")
    bins = defaultdict(list)
    routine = [r for r in rows if r["REPORT_TYPE"] == "FM-15"]
    for row in routine:
        stamp = datetime.fromisoformat(row["DATE"]).replace(tzinfo=timezone.utc)
        if stamp.year != YEAR:
            raise ValueError("Unexpected source year")
        bins[stamp.replace(minute=0, second=0)].append(row)
    start = datetime(YEAR, 1, 1, tzinfo=timezone.utc)
    all_hours = [start + timedelta(hours=i) for i in range(366 * 24)]
    work = [t for t in all_hours if t.astimezone(TZ).year == YEAR
            and t.astimezone(TZ).weekday() < 5 and 8 <= t.astimezone(TZ).hour < 17]
    selected = {t: choose(bins.get(t, [])) for t in all_hours}
    usable = {t: values(selected[t]) if selected[t] else None for t in all_hours}
    precipitation = precipitation_screen(bins, all_hours)
    monthly = []
    for month in range(1, 13):
        hours = [t for t in work if t.astimezone(TZ).month == month]
        for scenario, minimum in SCENARIOS.items():
            passed = [(t, qualifies(usable[t], minimum)) for t in hours]
            runs = long_runs(passed)
            combined = [(t, p and precipitation[t] == "reported_zero") for t, p in passed]
            rain_runs = long_runs(combined)
            monthly.append({
                "station": station, "label": label, "month": f"2024-{month:02}",
                "scenario": scenario, "expected_hours": len(hours),
                "no_report_hours": sum(not bins.get(t) for t in hours),
                "unusable_report_hours": sum(bool(bins.get(t)) and usable[t] is None for t in hours),
                "usable_hours": sum(usable[t] is not None for t in hours),
                "passing_hours": sum(p for _, p in passed),
                "runs_4plus": len(runs), "hours_in_runs_4plus": sum(map(len, runs)),
                "precip_reported_zero_hours": sum(precipitation[t] == "reported_zero" for t in hours),
                "precip_evidence_hours": sum(precipitation[t] == "precipitation_evidence" for t in hours),
                "precip_unknown_hours": sum(precipitation[t] == "unknown" for t in hours),
                "ambient_pass_precip_evidence_hours": sum(p and precipitation[t] == "precipitation_evidence" for t, p in passed),
                "ambient_pass_precip_unknown_hours": sum(p and precipitation[t] == "unknown" for t, p in passed),
                "precip_screen_passing_hours": sum(p for _, p in combined),
                "precip_screen_runs_4plus": len(rain_runs),
                "precip_screen_hours_in_runs_4plus": sum(map(len, rain_runs)),
            })
    audit = {
        "station": station, "label": label, "name": name, "latitude": lat, "longitude": lon,
        "url": f"https://noaa-global-hourly-pds.s3.amazonaws.com/{YEAR}/{station}.csv",
        "sha256": hashlib.sha256(blob).hexdigest(), "bytes": len(blob),
        "source_rows": len(rows), "first_source_utc": min(r["DATE"] for r in rows),
        "last_source_utc": max(r["DATE"] for r in rows),
        "routine_rows": len(routine), "expected_utc_hours": len(all_hours),
        "routine_hour_bins": len(bins),
        "multiple_routine_report_bins": sum(len(v) > 1 for v in bins.values()),
        "conflicting_selected_bins": sum(bool(bins.get(t)) and selected[t] is None for t in all_hours),
        "no_report_utc_hours": sum(not bins.get(t) for t in all_hours),
        "unusable_report_utc_hours": sum(bool(bins.get(t)) and usable[t] is None for t in all_hours),
        "usable_utc_hours": sum(v is not None for v in usable.values()),
        "routine_quality_flags": {
            key: dict(sorted(Counter(r[key].split(",")[-1] for r in routine).items()))
            for key in ("TMP", "DEW", "WND")
        },
        "precipitation_field_frequencies": {
            key: dict(sorted(Counter(r.get(key, "") or "(absent)" for r in routine).items()))
            for key in PRECIP_FIELDS if any(r.get(key) for r in routine)
        },
    }
    return monthly, audit


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--cache", type=Path, default=ROOT / ".cache/roof-weather")
    parser.add_argument("--refresh", action="store_true", help="Re-download; hashes reveal upstream revisions")
    parser.add_argument("--output", type=Path, default=OUT)
    args = parser.parse_args()
    args.cache.mkdir(parents=True, exist_ok=True)
    monthly, audits = [], []
    for station in STATIONS:
        path = args.cache / f"{station}.csv"
        if args.refresh or not path.exists():
            with urlopen(f"https://noaa-global-hourly-pds.s3.amazonaws.com/{YEAR}/{station}.csv", timeout=90) as res:
                blob = res.read()
            analyze(station, blob)  # Validate before caching.
            path.write_bytes(blob)
        records, audit = analyze(station, path.read_bytes())
        monthly.extend(records)
        audits.append(audit)
    for row in monthly:
        assert row["expected_hours"] == row["no_report_hours"] + row["unusable_report_hours"] + row["usable_hours"]
        assert row["hours_in_runs_4plus"] <= row["passing_hours"] <= row["usable_hours"]
        assert row["expected_hours"] == sum(row[k] for k in (
            "precip_reported_zero_hours", "precip_evidence_hours", "precip_unknown_hours"))
        assert row["passing_hours"] == sum(row[k] for k in (
            "precip_screen_passing_hours", "ambient_pass_precip_evidence_hours", "ambient_pass_precip_unknown_hours"))
        assert row["precip_screen_hours_in_runs_4plus"] <= row["hours_in_runs_4plus"]
        assert row["precip_screen_hours_in_runs_4plus"] <= row["precip_screen_passing_hours"]
    for audit in audits:
        assert audit["routine_hour_bins"] + audit["no_report_utc_hours"] == audit["expected_utc_hours"]
        assert audit["usable_utc_hours"] + audit["unusable_report_utc_hours"] == audit["routine_hour_bins"]
    output = {
        "year": YEAR, "timezone": "America/New_York",
        "screen": "Weekdays 08:00–16:59 local; 10 or 15 to 35 C air; air minus dew point >=3 C; wind <=4.5 m/s; QC 1/5 only. No precipitation screen.",
        "precipitation_screen": "Separate retrospective comparison: selected report must contain explicit 1-hour AA zero with QC 1/5 and condition 9; trace is precipitation. Any overlapping routine AA positive/trace accumulation vetoes the preceding hour; uncertain overlaps or supplied AU/AW reports remain unknown. No inferred roof dryness or installable days.",
        "stations": audits, "monthly": monthly,
    }
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "roof-weather-2024.json").write_text(json.dumps(output, indent=2) + "\n")
    with (args.output / "roof-weather-2024.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(monthly[0]))
        writer.writeheader()
        writer.writerows(monthly)
    (args.output / "roof-weather-analysis.py").write_bytes(Path(__file__).read_bytes())
    method = APP / "docs/roof-weather-method.md"
    if method.exists():
        (args.output / "roof-weather-method.md").write_bytes(method.read_bytes())
    for audit in audits:
        for scenario in SCENARIOS:
            subset = [r for r in monthly if r["station"] == audit["station"] and r["scenario"] == scenario]
            print(audit["label"], scenario, {k: sum(r[k] for r in subset) for k in list(subset[0])[4:]})


if __name__ == "__main__":
    main()