#!/usr/bin/env python3
"""
PHOENIX ranger_truth-gap replay — Fusion-GNN + Conformal prototype.

Honest design note: the shipped replay artifact contains EVENT/CANDIDATE-level records
(single_frame_candidates, firms_cluster_caution_shadow, priors, truth, ECMWF wind, DNB
granules) — NOT geostationary pixel frames (they rolled off cache; probe already proved
them zero-contrast). So this is a RELATIONAL GNN over event NODES, with wind-advection
edges, testing the real question: does multi-modal weak-signal fusion cross when every
single modality is sub-floor? Conformal wrapper bounds the false-negative rate.

Dependency-light on purpose: pure-torch message passing (no torch_geometric) + from-scratch
split/Mondrian conformal (no TorchCP) so it runs on a fresh pod without version friction.
(GraphFire-X directed-edge pattern arXiv 2512.20813; CRC-SGAD conformal-on-GNN arXiv 2504.02248.)

Run:  python train_fusiongnn.py --data ./replay --out ./verdict_card.json
"""
import argparse, json, math, os, glob, sys
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F

# ----------------------------- geo helpers -----------------------------------
R_EARTH_KM = 6371.0
ranger_truth = [(37.4602054, 13.2264283), (37.456807, 13.224735)]   # truth p1,p2
AOI_C = (37.4585, 13.2256)

def haversine_km(lat1, lon1, lat2, lon2):
    p = math.pi / 180
    a = (math.sin((lat2-lat1)*p/2)**2
         + math.cos(lat1*p)*math.cos(lat2*p)*math.sin((lon2-lon1)*p/2)**2)
    return 2*R_EARTH_KM*math.asin(math.sqrt(min(1.0, a)))

def bearing_deg(lat1, lon1, lat2, lon2):
    """Initial bearing from pt1 -> pt2, degrees clockwise from north."""
    p = math.pi/180
    y = math.sin((lon2-lon1)*p)*math.cos(lat2*p)
    x = (math.cos(lat1*p)*math.sin(lat2*p)
         - math.sin(lat1*p)*math.cos(lat2*p)*math.cos((lon2-lon1)*p))
    return (math.degrees(math.atan2(y, x)) + 360) % 360

# ----------------------------- data loading ----------------------------------
def _read(data_dir, name):
    """Load a db CSV defensively; return empty frame if absent."""
    for p in (os.path.join(data_dir, name),
              os.path.join(data_dir, "db", name),
              os.path.join(data_dir, "ranger_truth_20260630", "db", name)):
        if os.path.exists(p):
            try:
                return pd.read_csv(p)
            except Exception as e:
                print(f"warn: failed to read {p}: {e}", file=sys.stderr)
    print(f"warn: {name} not found under {data_dir}", file=sys.stderr)
    return pd.DataFrame()

def _col(df, *cands, default=np.nan):
    for c in cands:
        if c in df.columns:
            return df[c]
    return pd.Series([default]*len(df))

def dnb_features(data_dir):
    """Best-effort DNB night-glow radiance near the AOI, per granule.
    Reads VNP02DNB .nc if satpy/netCDF present; else returns None (pipeline still runs)."""
    granules = []
    for d in (data_dir, os.path.join(data_dir, "dnb"),
              os.path.join(data_dir, "ranger_truth_20260630", "dnb")):
        granules += glob.glob(os.path.join(d, "VNP02DNB*.nc"))
    granules = sorted(set(g for g in granules if not g.endswith(".processed")))
    if not granules:
        print("warn: no DNB granules found; DNB feature = 0", file=sys.stderr)
        return None
    try:
        from netCDF4 import Dataset  # noqa
    except Exception:
        print("warn: netCDF4 unavailable; DNB feature = 0 (install for DNB signal)", file=sys.stderr)
        return None
    vals = []
    for g in granules:
        try:
            ds = Dataset(g)
            # VNP02DNB radiance is under observation_data/DNB_observations (product-dependent).
            grp = ds.groups.get("observation_data")
            if grp and "DNB_observations" in grp.variables:
                arr = np.asarray(grp.variables["DNB_observations"][:], dtype="float32")
                vals.append(float(np.nanpercentile(arr, 99)))  # bright-tail proxy
            ds.close()
        except Exception as e:
            print(f"warn: DNB read {os.path.basename(g)}: {e}", file=sys.stderr)
    if not vals:
        return None
    return {"dnb_p99_max": float(np.nanmax(vals)), "dnb_n_granules": len(granules)}

def wind_uv_at(ecmwf, lat, lon):
    """Nearest-grid ECMWF 10m wind vector (u10=east, v10=north). Calm fallback."""
    if ecmwf.empty:
        return 0.0, 0.0
    la = _col(ecmwf, "lat_grid", "lat").to_numpy(dtype="float64")
    lo = _col(ecmwf, "lon_grid", "lon").to_numpy(dtype="float64")
    u = _col(ecmwf, "u10", default=0.0).to_numpy(dtype="float64")
    v = _col(ecmwf, "v10", default=0.0).to_numpy(dtype="float64")
    d2 = (la-lat)**2 + (lo-lon)**2
    i = int(np.nanargmin(d2))
    return float(u[i]), float(v[i])

def wind_at(ecmwf, lat, lon):
    """(speed, from-direction-deg) derived from the wind vector."""
    uu, vv = wind_uv_at(ecmwf, lat, lon)
    spd = math.hypot(uu, vv)
    frm = (math.degrees(math.atan2(-uu, -vv)) + 360) % 360
    return spd, frm

# ----------------------------- graph build -----------------------------------
def build_graph(data_dir, k_spatial=6, wind_link_km=12.0):
    truth   = _read(data_dir, "truth.csv")
    firms_c = _read(data_dir, "marginal_controls_firmscaution.csv")
    sfc     = _read(data_dir, "marginal_controls_singleframe.csv")
    priors  = _read(data_dir, "firms_anchored_priors.csv")
    recur   = _read(data_dir, "external_fires_recent.csv")
    corrob  = _read(data_dir, "corroboration_signals.csv")
    citizen = _read(data_dir, "citizen_reports.csv")
    ecmwf   = _read(data_dir, "ecmwf_hres_window.csv")
    dnb     = dnb_features(data_dir)
    dnb_p99 = (dnb or {}).get("dnb_p99_max", 0.0)

    nodes = []   # each: dict(lat, lon, label, feat..., group)
    def add(lat, lon, label, kind, **f):
        if not np.isfinite(lat) or not np.isfinite(lon):
            return
        nodes.append(dict(lat=float(lat), lon=float(lon), label=int(label), kind=kind, **f))

    # POSITIVES: firms caution clusters that were real-but-subsurface (own_confirmed/is_recurring)
    for _, r in firms_c.iterrows():
        add(_g(r,"lat"), _g(r,"lon"), 1, "firms_caution",
            frp=_g(r,"max_frp_w",0.0), n_sensors=_g(r,"n_sensors",0.0),
            n_passes=_g(r,"n_passes",0.0), recurring=_g(r,"is_recurring",0.0),
            weather_clear=_g(r,"weather_clear",0.0))
    # CANDIDATES (weak, unpromoted single-frame): treat as hard negatives/uncertain (label 0)
    for _, r in sfc.iterrows():
        add(_g(r,"lat"), _g(r,"lon"), 0, "single_frame",
            t_fire_k=_g(r,"t_fire_k",0.0), frp=_g(r,"frp_mw",0.0),
            persistence=_g(r,"persistence_seen_n",1.0))
    # TRUTH: the ranger_truth event (held-out target, positive)
    for _, r in truth.iterrows():
        add(_g(r,"lat"), _g(r,"lon"), 1, "ranger_truth_truth")

    if not nodes:
        raise SystemExit("no nodes built — check CSV contents")

    df = pd.DataFrame(nodes)
    lat = df["lat"].to_numpy(); lon = df["lon"].to_numpy(); N = len(df)

    # ---- per-node fused weak-signal features (the whole point) ----
    def near_count(tbl, latc="lat", lonc="lng", radius=8.0):
        if tbl.empty: return np.zeros(N)
        tl = _col(tbl, latc, "lat").to_numpy(dtype="float64")
        to = _col(tbl, lonc, "lng", "lon").to_numpy(dtype="float64")
        out = np.zeros(N)
        for i in range(N):
            d = np.array([haversine_km(lat[i],lon[i],a,b) for a,b in zip(tl,to)])
            out[i] = int(np.sum(d <= radius))
        return out
    recur_n  = near_count(recur, "lat", "lng", 10.0)      # recurrence-anchored prior
    corrob_n = near_count(corrob, "lat", "lng", 8.0)      # multi-source corroboration
    prior_n  = near_count(priors, "lat", "lon", 7.0)      # firms/citizen anchored prior
    citiz_n  = near_count(citizen, "lat", "lon", 5.0)     # citizen report proximity

    feat = np.stack([
        _fnum(df,"frp"), _fnum(df,"t_fire_k"), _fnum(df,"persistence"),
        _fnum(df,"n_sensors"), _fnum(df,"n_passes"), _fnum(df,"recurring"),
        _fnum(df,"weather_clear"),
        recur_n, corrob_n, prior_n, citiz_n,
        np.full(N, dnb_p99),                              # DNB night-glow bright-tail
    ], axis=1).astype("float32")
    # standardize
    mu = np.nanmean(feat,0); sd = np.nanstd(feat,0)+1e-6
    feat = np.nan_to_num((feat-mu)/sd)

    # ---- edges: spatial kNN (undirected) + wind-advection (directed upwind->downwind) ----
    src, dst, ew = [], [], []
    for i in range(N):
        d = np.array([haversine_km(lat[i],lon[i],lat[j],lon[j]) if j!=i else 1e9 for j in range(N)])
        for j in np.argsort(d)[:k_spatial]:
            src += [i, j]; dst += [j, i]; ew += [[1.0, 0.0], [1.0, 0.0]]  # [spatial,wind]
        spd, frm = wind_at(ecmwf, lat[i], lon[i])
        if spd > 0.5:
            for j in range(N):
                if j==i: continue
                dij = haversine_km(lat[i],lon[i],lat[j],lon[j])
                if dij > wind_link_km: continue
                brg = bearing_deg(lat[i],lon[i],lat[j],lon[j])   # i -> j direction
                downwind = (frm + 180) % 360                     # wind blows toward here
                align = math.cos(math.radians(brg - downwind))   # 1 = j is straight downwind of i
                if align > 0.3:
                    w = align * min(spd,15)/15.0
                    src += [i]; dst += [j]; ew += [[0.0, float(w)]]  # directed advection edge
    edge_index = torch.tensor([src, dst], dtype=torch.long)
    edge_attr  = torch.tensor(ew, dtype=torch.float32)

    # per-node wind vector (u,v) + lat/lon — for the ADGNO advection operator
    node_uv = np.array([wind_uv_at(ecmwf, lat[i], lon[i]) for i in range(N)], dtype="float32")
    node_ll = np.stack([lat, lon], axis=1).astype("float32")

    x = torch.tensor(feat, dtype=torch.float32)
    y = torch.tensor(df["label"].to_numpy(), dtype=torch.float32)
    is_truth = torch.tensor((df["kind"]=="ranger_truth_truth").to_numpy(), dtype=torch.bool)
    # Mondrian region id: coarse 0.5deg lat/lon bins for per-region conformal
    region = ((df["lat"]//0.5).astype(int).astype(str)+"_"+(df["lon"]//0.5).astype(int).astype(str)).to_numpy()
    return dict(x=x, y=y, edge_index=edge_index, edge_attr=edge_attr,
                node_uv=torch.tensor(node_uv), node_ll=torch.tensor(node_ll),
                is_truth=is_truth, region=region, df=df,
                meta=dict(n_nodes=N, n_pos=int(y.sum()), dnb=dnb or {}, feat_dim=feat.shape[1]))

def _g(r, c, default=np.nan):
    return r[c] if c in r and pd.notna(r[c]) else default
def _fnum(df, c):
    return pd.to_numeric(df[c], errors="coerce").to_numpy(dtype="float64") if c in df else np.zeros(len(df))

# ----------------------------- model (pure torch) ----------------------------
class WindGAT(nn.Module):
    """Attention message passing with edge features [spatial, wind]. No torch_geometric."""
    def __init__(self, in_dim, hid=32, layers=2):
        super().__init__()
        self.lin_in = nn.Linear(in_dim, hid)
        self.att = nn.ModuleList([nn.Linear(2*hid+2, 1) for _ in range(layers)])
        self.upd = nn.ModuleList([nn.Linear(2*hid, hid) for _ in range(layers)])
        self.out = nn.Linear(hid, 1)
        self.layers = layers
    def forward(self, x, edge_index, edge_attr):
        h = F.relu(self.lin_in(x))
        s, d = edge_index
        for l in range(self.layers):
            m_in = torch.cat([h[s], h[d], edge_attr], dim=1)
            a = torch.sigmoid(self.att[l](m_in)).squeeze(1)         # edge gate
            msg = a.unsqueeze(1) * h[s]
            agg = torch.zeros_like(h).index_add_(0, d, msg)
            deg = torch.zeros(h.size(0), device=h.device).index_add_(0, d, a).clamp(min=1e-6)
            agg = agg / deg.unsqueeze(1)
            h = F.relu(self.upd[l](torch.cat([h, agg], dim=1))) + h
        return self.out(h).squeeze(1)   # logits

# ----------------------------- conformal (from scratch) ----------------------
def split_conformal_fnr(scores_pos_cal, alpha=0.05):
    """CRC-style: pick a score threshold t so calibrated FNR (missed positives) <= alpha.
    Nonconformity for a positive = -score (higher score => more 'fire')."""
    s = np.sort(np.asarray(scores_pos_cal))          # ascending
    if len(s)==0: return -np.inf
    # allow at most alpha fraction of positives below threshold -> quantile at alpha
    q = np.floor(alpha*(len(s)+1))/len(s)
    q = min(max(q, 0.0), 1.0)
    return float(np.quantile(s, q, method="lower"))

def mondrian_thresholds(scores, y, region, alpha=0.05):
    th = {}
    for reg in np.unique(region):
        m = (region==reg) & (y==1)
        th[reg] = split_conformal_fnr(scores[m], alpha) if m.sum()>=5 else None
    global_t = split_conformal_fnr(scores[y==1], alpha)
    return th, global_t

# ----------------------------- one trial (single seed) -----------------------
def run_trial(G, args, seed):
    torch.manual_seed(seed); np.random.seed(seed)
    x, y = G["x"], G["y"]; ei, ea = G["edge_index"], G["edge_attr"]
    if args.no_citizen:
        x = x.clone(); x[:, 10] = 0.0   # zero ONLY citiz_n (idx 10); keep prior_n (9)

    if args.model == "adgno":
        from adgno import ADGNO
        keep = ea[:,1]==0
        ei_a, ea_a = ei[:,keep], ea[keep]
        adv_w, deg = ADGNO.build_advection(ei_a, ea_a, G["node_uv"], G["node_ll"])
        if args.no_wind: adv_w = torch.zeros_like(adv_w)
        model = ADGNO(x.size(1)); fwd = lambda: model(x, ei_a, adv_w, deg)
    else:
        if args.no_wind:
            keep = ea[:,1]==0; ei, ea = ei[:,keep], ea[keep]
        model = WindGAT(x.size(1)); fwd = lambda: model(x, ei, ea)

    N = x.size(0); idx = np.arange(N)
    truth_mask = G["is_truth"].numpy()
    pool = idx[~truth_mask]
    rng = np.random.default_rng(seed); rng.shuffle(pool)
    n_cal = max(10, int(0.3*len(pool)))
    cal_idx, tr_idx = pool[:n_cal], pool[n_cal:]        # calib held out from training
    test_idx = idx[truth_mask]

    opt = torch.optim.Adam(model.parameters(), lr=1e-2, weight_decay=1e-4)
    pos_w = torch.tensor([(y[tr_idx]==0).sum()/max(1,(y[tr_idx]==1).sum())])
    tr = torch.tensor(tr_idx)
    for _ in range(args.epochs):
        model.train(); opt.zero_grad()
        loss = F.binary_cross_entropy_with_logits(fwd()[tr], y[tr], pos_weight=pos_w)
        loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        scores = torch.sigmoid(fwd()).numpy()

    # conformal calibrated ONLY on held-out calibration positives (no leakage)
    yc = y.numpy()
    cal_pos = np.array([i for i in cal_idx if yc[i]==1])
    th_glob = split_conformal_fnr(scores[cal_pos], args.alpha)
    caught = [bool(scores[ti] >= th_glob) for ti in test_idx]
    footprint = float(np.mean(scores >= th_glob))
    return dict(caught_any=any(caught), ranger_truth_scores=[float(scores[i]) for i in test_idx],
                footprint=footprint, threshold=float(th_glob))

# ----------------------------- multi-seed driver -----------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="./replay")
    ap.add_argument("--out", default="./verdict_card.json")
    ap.add_argument("--alpha", type=float, default=0.05)
    ap.add_argument("--epochs", type=int, default=300)
    ap.add_argument("--seeds", type=int, default=7, help="repeats for robustness (mean +/- std)")
    ap.add_argument("--model", choices=["gat","adgno"], default="gat")
    ap.add_argument("--no-wind", action="store_true")
    ap.add_argument("--no-citizen", action="store_true")
    args = ap.parse_args()

    G = build_graph(args.data)          # deterministic graph, built once
    trials = [run_trial(G, args, s) for s in range(args.seeds)]

    caught_rate = float(np.mean([t["caught_any"] for t in trials]))     # frac of seeds that caught
    foot = np.array([t["footprint"] for t in trials])
    gsc  = np.array([t["ranger_truth_scores"] for t in trials])              # [seeds, n_truth]
    verdict = dict(
        question="Would fusion-GNN + conformal catch the ranger_truth clear-sky sub-threshold fire?",
        model=args.model, seeds=args.seeds, alpha=args.alpha,
        n_nodes=G["meta"]["n_nodes"], n_positives=G["meta"]["n_pos"],
        feat_dim=G["meta"]["feat_dim"], dnb=G["meta"]["dnb"],
        caught_rate=round(caught_rate,3),                   # 1.0 = caught on every seed
        caught_any=bool(caught_rate > 0.5),                 # majority verdict
        ranger_truth_score_mean=[round(float(m),4) for m in gsc.mean(0)],
        ranger_truth_score_std =[round(float(s),4) for s in gsc.std(0)],
        alert_footprint_mean=round(float(foot.mean()),4),
        alert_footprint_std =round(float(foot.std()),4),
        ablation=dict(no_wind=args.no_wind, no_citizen=args.no_citizen),
        note=("Event-level relational fusion (no geostationary frames — rolled off; probe proved "
              "thermal null). Positive=real-but-subsurface fires + ranger_truth truth. Multi-seed: "
              "caught_rate is the fraction of seeds whose conformal-calibrated score cleared the "
              "FNR bound at a truth node — robust to single-seed flukes."),
    )
    with open(args.out, "w") as f: json.dump(verdict, f, indent=2)
    print(json.dumps(verdict, indent=2))

if __name__ == "__main__":
    main()