"""Bosch GNN: each product is a subgraph of the station network.

The tabular model sees a product as 256 numbers with no notion that station
L3_S29 feeds L3_S30. This model sees the actual routing: nodes are the
stations a product visited, edges are the transitions it made, node features
are that station's measurements for *that* product. Message passing can
therefore combine a reading with its neighbours' readings along the route.

Comparison rules are the same as run_model.py — same temporal split, same
official metrics — because the question is whether the structure buys
anything, not whether a GNN can be trained.

Run: python 03_bosch/run_gnn.py   (needs prepare.py + build_graph.py)
"""
from __future__ import annotations

import sys
import time
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
from torch_geometric.nn import GraphConv, global_max_pool, global_mean_pool

HERE = Path(__file__).parent
sys.path.insert(0, str(HERE.parent))
sys.path.insert(0, str(HERE))

from common import metrics  # noqa: E402

CACHE = HERE / "cache"
FIGS = HERE / "figures"
OUT = HERE / "outputs"

COST_FN, COST_FP = 50.0, 1.0
NODE_AGGS = ["n", "mean", "min", "max", "std"]
EPOCHS = 12
BATCH = 512
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1.18M products is far more than a CPU GNN needs to learn this; subsample the
# overwhelming negative class and keep every failure.
NEG_PER_POS = 25


def load_tables():
    times = pd.read_parquet(CACHE / "station_times.parquet")
    num = pd.read_parquet(CACHE / "station_numeric.parquet")
    assert (times["Id"].to_numpy() == num["Id"].to_numpy()).all()
    stations = [c for c in times.columns if c != "Id"]
    return times, num, stations


def build_split(times, num, stations):
    """Same temporal split as the tabular model, so results are comparable."""
    T = times[stations].to_numpy(dtype=np.float32)
    visited = ~np.isnan(T)
    y = num["Response"].to_numpy().astype(np.int64)
    with np.errstate(all="ignore"):
        start = np.where(visited.any(axis=1),
                         np.nanmin(np.nan_to_num(T, nan=np.inf), axis=1), np.nan)
    known = ~np.isnan(start)
    order = np.argsort(np.where(known, start, np.inf), kind="stable")[:known.sum()]
    n = len(order)
    return {
        "fit": order[:int(0.55 * n)],
        "val": order[int(0.55 * n):int(0.70 * n)],
        "test": order[int(0.70 * n):],
    }, T, visited, y


def subsample(idx: np.ndarray, y: np.ndarray, seed: int = 0) -> np.ndarray:
    """Keep all failures, sample NEG_PER_POS negatives per failure."""
    rng = np.random.default_rng(seed)
    pos = idx[y[idx] == 1]
    neg = idx[y[idx] == 0]
    take = min(len(neg), NEG_PER_POS * max(1, len(pos)))
    return np.concatenate([pos, rng.choice(neg, take, replace=False)])


def make_graphs(idx, T, visited, node_feats, y, stations) -> list[Data]:
    """One PyG Data per product: its visited stations as an ordered chain."""
    n_st = len(stations)
    st_index = torch.arange(n_st)
    out = []
    for i in idx:
        vis = np.flatnonzero(visited[i])
        if len(vis) == 0:
            continue
        vis = vis[np.argsort(T[i, vis])]          # station order along the route
        x = torch.tensor(node_feats[i][vis], dtype=torch.float32)
        # station identity as a learnable embedding index
        sid = st_index[torch.tensor(vis.copy(), dtype=torch.long)]
        if len(vis) > 1:
            src = torch.arange(len(vis) - 1)
            ei = torch.stack([torch.cat([src, src + 1]),
                              torch.cat([src + 1, src])])   # undirected chain
        else:
            ei = torch.zeros((2, 0), dtype=torch.long)
        out.append(Data(x=x, edge_index=ei, station=sid,
                        y=torch.tensor([y[i]], dtype=torch.float32)))
    return out


class ProductGNN(nn.Module):
    def __init__(self, n_feat: int, n_stations: int, hidden: int = 64,
                 emb: int = 16):
        super().__init__()
        self.emb = nn.Embedding(n_stations, emb)
        self.inp = nn.Linear(n_feat + emb, hidden)
        self.c1 = GraphConv(hidden, hidden)
        self.c2 = GraphConv(hidden, hidden)
        self.norm = nn.BatchNorm1d(hidden)
        self.head = nn.Sequential(nn.Linear(2 * hidden, 64), nn.ReLU(),
                                  nn.Dropout(0.2), nn.Linear(64, 1))

    def forward(self, d):
        h = torch.cat([d.x, self.emb(d.station)], dim=1)
        h = torch.relu(self.inp(h))
        h = torch.relu(self.c1(h, d.edge_index))
        h = self.norm(torch.relu(self.c2(h, d.edge_index)) + h)
        pooled = torch.cat([global_mean_pool(h, d.batch),
                            global_max_pool(h, d.batch)], dim=1)
        return self.head(pooled).squeeze(1)


@torch.no_grad()
def predict(model, loader) -> tuple[np.ndarray, np.ndarray]:
    model.eval()
    ps, ys = [], []
    for b in loader:
        b = b.to(DEVICE)
        ps.append(torch.sigmoid(model(b)).cpu().numpy())
        ys.append(b.y.cpu().numpy())
    return np.concatenate(ps), np.concatenate(ys)


def main():
    FIGS.mkdir(exist_ok=True)
    OUT.mkdir(exist_ok=True)
    torch.manual_seed(0)

    times, num, stations = load_tables()
    split, T, visited, y = build_split(times, num, stations)
    print(f"device {DEVICE}; {len(y):,} products, fail rate {y.mean():.4%}")

    # node feature tensor: [n_products, n_stations, n_aggs].
    # Two stations (e.g. L3_S42) appear in the date file but take no numeric
    # measurements. They are still real nodes on the route, so they keep a
    # node with zero features and are identified by their station embedding.
    n_rows = len(num)
    blocks = []
    no_measurements = []
    for s in stations:
        cols = [f"{s}__{a}" for a in NODE_AGGS]
        if all(c in num.columns for c in cols):
            blocks.append(num[cols].to_numpy(dtype=np.float32))
        else:
            no_measurements.append(s)
            blocks.append(np.zeros((n_rows, len(NODE_AGGS)), dtype=np.float32))
    if no_measurements:
        print(f"stations with routing but no measurements: {no_measurements}")
    node_feats = np.stack(blocks, axis=1)
    node_feats = np.nan_to_num(node_feats, nan=0.0, posinf=0.0, neginf=0.0)
    # standardise per (station, agg) using the fit window only
    fit_idx = split["fit"]
    mu = node_feats[fit_idx].mean(axis=0, keepdims=True)
    sd = node_feats[fit_idx].std(axis=0, keepdims=True)
    sd[sd < 1e-6] = 1.0
    node_feats = (node_feats - mu) / sd
    print(f"node features {node_feats.shape}")

    # Training is subsampled (1.18M products is far more than a CPU GNN needs),
    # but validation and test are kept at FULL prevalence. Thresholds must be
    # chosen at the prevalence they will face, and models A/B chose theirs on
    # the full validation window — a subsampled one would not be comparable.
    fit_s = subsample(split["fit"], y, seed=0)
    val_s = split["val"]
    print(f"training graphs {len(fit_s):,} (fail rate {y[fit_s].mean():.2%}, "
          f"subsampled), val {len(val_s):,} and test {len(split['test']):,} "
          f"(full prevalence {y[val_s].mean():.3%} / {y[split['test']].mean():.3%})")

    t0 = time.time()
    g_fit = make_graphs(fit_s, T, visited, node_feats, y, stations)
    g_val = make_graphs(val_s, T, visited, node_feats, y, stations)
    g_test = make_graphs(split["test"], T, visited, node_feats, y, stations)
    print(f"graphs built in {time.time() - t0:.0f}s", flush=True)

    dl_fit = DataLoader(g_fit, batch_size=BATCH, shuffle=True)
    dl_val = DataLoader(g_val, batch_size=2048)
    dl_test = DataLoader(g_test, batch_size=2048)

    model = ProductGNN(len(NODE_AGGS), len(stations)).to(DEVICE)
    pos_weight = torch.tensor([(y[fit_s] == 0).sum() / max(1, (y[fit_s] == 1).sum())],
                              device=DEVICE)
    loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=EPOCHS)

    best_ap, best_state = -1.0, None
    for epoch in range(EPOCHS):
        model.train()
        t0 = time.time()
        for b in dl_fit:
            b = b.to(DEVICE)
            opt.zero_grad()
            loss_fn(model(b), b.y).backward()
            nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
        sched.step()
        vp, vy = predict(model, dl_val)
        ap = metrics.pr_auc(vy, vp)
        print(f"  epoch {epoch + 1}/{EPOCHS}: val PR-AUC {ap:.4f} "
              f"({time.time() - t0:.0f}s)", flush=True)
        if ap > best_ap:
            best_ap = ap
            best_state = {k: v.clone() for k, v in model.state_dict().items()}
    model.load_state_dict(best_state)

    vp, vy = predict(model, dl_val)
    thr_mcc, _ = metrics.best_threshold_mcc(vy, vp)
    thr_cost, _ = metrics.best_threshold(vy, vp, COST_FN, COST_FP)
    tp_, ty_ = predict(model, dl_test)
    _, mcc_oracle = metrics.best_threshold_mcc(ty_, tp_)

    rec = {
        "model": "C: GNN on product routing subgraphs",
        "mcc@val_threshold": metrics.mcc(ty_, (tp_ >= thr_mcc).astype(int)),
        "mcc@oracle_threshold": mcc_oracle,
        "pr_auc": metrics.pr_auc(ty_, tp_),
        "capture@20%": metrics.decile_capture(ty_, tp_, 0.2),
        "capture@5%": metrics.decile_capture(ty_, tp_, 0.05),
        "recall@1%fpr": metrics.recall_at_fpr(ty_, tp_, 0.01),
        "cost/part": metrics.expected_cost(ty_, (tp_ >= thr_cost).astype(int),
                                           COST_FN, COST_FP),
        "n_features": f"{len(NODE_AGGS)}/station + routing",
    }
    tab = pd.read_csv(OUT / "model_comparison.csv")
    combined = pd.concat([tab, pd.DataFrame([rec])], ignore_index=True).set_index("model")
    combined.to_csv(OUT / "model_comparison_with_gnn.csv")
    np.savez_compressed(CACHE / "gnn_test_scores.npz", y=ty_, score=tp_)
    print("\n=== XGBoost tracks vs GNN, same split and metrics ===")
    print(combined.round(4).to_string())

    # --- capture curve against the tabular models
    z = np.load(CACHE / "test_scores.npz")
    fracs = np.linspace(0.01, 1.0, 60)
    fig, ax = plt.subplots(figsize=(7.5, 5.5))
    for key, label, color in [("m0", "A: station aggregates (XGBoost)", "steelblue"),
                              ("m1", "B: + path/graph features (XGBoost)", "darkorange")]:
        caps = [metrics.decile_capture(z["y"], z[key], f) for f in fracs]
        ax.plot(fracs, caps, label=label, color=color)
    ax.plot(fracs, [metrics.decile_capture(ty_, tp_, f) for f in fracs],
            label="C: GNN on routing subgraphs", color="seagreen")
    ax.plot(fracs, fracs, ls=":", color="gray", label="random inspection")
    ax.set(xlabel="fraction of products inspected (highest risk first)",
           ylabel="fraction of all failures caught",
           title="Does message passing over the route beat flat features?\n"
                 "(test period, identical temporal split)")
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(FIGS / "gnn_vs_tabular.png", dpi=150)
    plt.close(fig)


if __name__ == "__main__":
    main()