"""pull_azmet.py

Freezes the AZMet station year used by Part 0-B and 0-C, and computes the
check values those exercises ask students to reproduce.

The Arizona Meteorological Network publishes daily observations with no key
and no account. Station metadata, including elevation, comes from the
published metadata table. Nothing here needs a GIS package.

    python pull_azmet.py
"""
from __future__ import annotations

import hashlib
import io
import json
import math
import urllib.request
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent.parent
SNAP = ROOT / "snapshots"
API = "https://api.azmet.arizona.edu/v1/observations/daily"
META = "https://azmet.arizona.edu/about/station-metadata"
YEAR = 2024
INTERVAL = "P365D"        # P365D spans exactly the calendar year
NUM = ["temp_air_maxC", "temp_air_meanC", "temp_air_minC",
       "precip_total_mm", "wind_vector_dir", "eto_azmet"]


def stations() -> pd.DataFrame:
    with urllib.request.urlopen(META, timeout=180) as resp:
        html = resp.read().decode("utf-8", errors="replace")
    t = pd.read_html(io.StringIO(html))[0]
    t.columns = ["name", "symbol", "station", "county", "lat", "lon",
                 "elev_ft", "elev_m", "status", "description"]
    return t


def daily(sid: str):
    url = f"{API}/{sid}/{YEAR}-01-01T00:00/{INTERVAL}"
    try:
        with urllib.request.urlopen(url, timeout=180) as resp:
            rows = json.loads(resp.read()).get("data", [])
    except Exception:
        return None
    if not rows:
        return None
    df = pd.DataFrame(rows)
    for c in NUM:
        if c in df.columns:
            df[c] = pd.to_numeric(df[c], errors="coerce")
            # AZMet encodes missing observations as -9999. Averaging a year
            # that contains two of them moves an annual mean by fifty
            # degrees, and nothing in the response marks the value as
            # absent. Mask sentinels before any arithmetic touches them.
            df.loc[df[c] <= -900, c] = np.nan
            df.loc[df[c] >= 9000, c] = np.nan
    return df


def circular_mean(series) -> float:
    a = np.deg2rad(pd.to_numeric(series, errors="coerce").dropna().to_numpy())
    if not len(a):
        return float("nan")
    return float(np.rad2deg(math.atan2(np.sin(a).mean(),
                                       np.cos(a).mean())) % 360)


def main():
    SNAP.mkdir(parents=True, exist_ok=True)
    st = stations()
    print(f"stations listed: {len(st)}  active: {(st.status=='Active').sum()}")
    rows = []
    for sid in sorted(st.station.unique()):
        df = daily(sid)
        if df is None or len(df) < 360:
            print(f"  skip {sid} ({0 if df is None else len(df)} rows)")
            continue
        rows.append({
            "station": sid, "n_days": len(df),
            "tmean_C": df.temp_air_meanC.mean(),
            "tmax_C": df.temp_air_maxC.mean(),
            "precip_mm": df.precip_total_mm.sum(),
            "zero_rain_days": int((df.precip_total_mm == 0).sum()),
            "eto_mm": df.eto_azmet.sum() if "eto_azmet" in df else np.nan,
            "wind_arith": df.wind_vector_dir.mean(),
            "wind_circ": circular_mean(df.wind_vector_dir),
        })
        print(f"  {sid} ok ({len(df)} rows)", flush=True)
    ann = pd.DataFrame(rows).merge(
        st[["station", "name", "county", "lat", "lon", "elev_m", "status"]],
        on="station", how="left")
    out = SNAP / f"azmet_{YEAR}_annual.csv"
    ann.to_csv(out, index=False)
    digest = hashlib.sha256(out.read_bytes()).hexdigest()
    print(f"\nfrozen: {out}\nrows {len(ann)}  sha256 {digest[:32]}")

    # ---- Part 0-C check values: elevation lapse rate + LOOCV ----------
    d = ann.dropna(subset=["tmean_C", "elev_m"])
    x, y = d.elev_m.to_numpy(float), d.tmean_C.to_numpy(float)
    slope, intercept = np.polyfit(x, y, 1)
    r = float(np.corrcoef(x, y)[0, 1])
    resid = []
    for i in range(len(x)):
        m = np.ones(len(x), bool)
        m[i] = False
        s2, i2 = np.polyfit(x[m], y[m], 1)
        resid.append(y[i] - (s2 * x[i] + i2))
    resid = np.array(resid)
    res = {
        "stations_full_year": int(len(ann)),
        "stations_in_regression": int(len(d)),
        "lapse_C_per_1000m": round(-slope * 1000, 3),
        "intercept_C": round(float(intercept), 3),
        "pearson_r": round(r, 4), "r_squared": round(r * r, 4),
        "loocv_rmse_C": round(float(np.sqrt((resid ** 2).mean())), 3),
        "loocv_mae_C": round(float(np.abs(resid).mean()), 3),
        "largest_residual_station": d.iloc[int(np.abs(resid).argmax())]["name"],
        "largest_residual_C": round(float(resid[np.abs(resid).argmax()]), 2),
        "elev_min_m": int(d.elev_m.min()), "elev_max_m": int(d.elev_m.max()),
        "sha256": digest,
    }
    (SNAP / "azmet_checkvalues.json").write_text(json.dumps(res, indent=2))
    print(json.dumps(res, indent=2))


if __name__ == "__main__":
    main()