"""Graph-free LSTM and topology-aware GNN for day-ahead probabilistic heat load.

Shared multi-output model over all 27 consumers. Each consumer's last 7 days
of load (+ observation mask + exogenous features) are encoded by a shared
LSTM; the GNN variant then runs 2 message-passing rounds over the 42-node
grid graph (15 junctions + 27 consumers, from Memmingen_pressure.yaml) before
decoding. Heads per consumer and step: 10 quantiles (pinball-trained) and a
Johnson SU distribution (NLL-trained). Losses are masked to observed targets.

Usage: python train_seq.py --model {lstm,gnn} [--smoke]
Pushes results/{model}.json and models/{model}.pt to
jantheman/heating-grid-forecast-benchmark.
"""

import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path

import numpy as np
import torch

REPO_DATA = "jantheman/heating-grid-forecast-data"
REPO_CODE = "jantheman/heating-grid-forecast-benchmark"
LEVELS = np.array([0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65, 0.75, 0.85, 0.95])
HORIZON = 96
SEQ = 7 * 96  # 672 history slots
HIDDEN = 128
EPOCHS = 40
BATCH = 16
N_VAL_ORIGINS = 30
SEED = 0


def download_all(root: Path) -> dict:
    from huggingface_hub import hf_hub_download
    (root / "heatingforecast").mkdir(exist_ok=True)
    for f in ["__init__.py", "data.py", "distributions.py", "metrics.py", "graph.py"]:
        hf_hub_download(REPO_CODE, f"heatingforecast/{f}", repo_type="model",
                        local_dir=root)
    files = {}
    for fn in ["epronet_full_20260427.csv", "weather/weather_memmingen_hourly.csv",
               "weather/holidays_bavaria.csv", "Memmingen_pressure.yaml",
               "graph/consumer_mapping.csv"]:
        files[fn] = hf_hub_download(REPO_DATA, fn, repo_type="dataset", local_dir=root)
    return files


def prepare(files: dict):
    import pandas as pd
    sys.path.insert(0, str(Path.cwd()))  # stdin-run safe: no __file__
    import heatingforecast.data as hfdata
    import heatingforecast.graph as hfgraph

    wide, _ = hfdata.load_wide(files["epronet_full_20260427.csv"])
    wide, obs_mask = hfdata.bridge_gaps(wide)
    feats = hfdata.exog_features(files["weather/weather_memmingen_hourly.csv"],
                                 files["weather/holidays_bavaria.csv"], wide.index)
    tr, va, te = hfdata.splits(wide)
    fit = tr | va

    Y = wide.values.astype(np.float32)              # (T, 27) kWh per 15-min
    M = obs_mask.values.astype(np.float32)
    X = feats.values.astype(np.float32)             # (T, F)

    # z-score exog on the fit window: raw scales (solar ~0-1000 W/m2, humidity
    # 0-100) blow up LSTM activations over 768 steps -> NaN. Trees don't care;
    # recurrent nets absolutely do.
    Xok = fit[:, None] & ~np.isnan(X)
    xmu = np.nanmean(np.where(Xok, X, np.nan), axis=0)
    xsd = np.nanstd(np.where(Xok, X, np.nan), axis=0) + 1e-6
    X = ((X - xmu) / xsd).astype(np.float32)

    # per-consumer standardization from the fit window (NaN-aware)
    ok = fit[:, None] & ~np.isnan(Y)
    mu = np.nanmean(np.where(ok, Y, np.nan), axis=0)
    sd = np.nanstd(np.where(ok, Y, np.nan), axis=0) + 1e-6
    Ys = (Y - mu) / sd

    g = hfgraph.load_graph(files["Memmingen_pressure.yaml"],
                           files["graph/consumer_mapping.csv"])
    assert g["consumers"] == list(wide.columns)

    # dense adjacency with pipe-length decay on consumer<->junction edges
    A = np.zeros((g["n_nodes"], g["n_nodes"]), dtype=np.float32)
    ei, ea = g["edge_index"], g["edge_attr"]
    for k in range(ei.shape[1]):
        u, v, L = ei[0, k], ei[1, k], ea[k, 0]
        A[u, v] = 1.0 / (1.0 + L / 500.0)
    d = A.sum(1, keepdims=True)
    A_norm = A / np.maximum(d, 1e-6)

    idx = {c: i for i, c in enumerate(g["consumers"])}
    consumer_nodes = np.array([idx[c] for c in g["consumers"]])   # node id per column

    return dict(Y=Ys, Yraw=Y, M=M, X=X, fit=fit, test=te,
                mu=mu.astype(np.float32), sd=sd.astype(np.float32),
                A=A_norm.astype(np.float32), consumer_nodes=consumer_nodes)


def daily_origins(index, day_mask):
    T = len(index)
    out = []
    for o in range(SEQ, T - 2 * HORIZON):
        if index[o].hour or index[o].minute:
            continue
        if day_mask[np.arange(o + HORIZON, o + 2 * HORIZON)].all():
            out.append(o)
    return np.asarray(out)


def make_batches(origins, D):
    Y, M, X = D["Y"], D["M"], D["X"]
    for i in range(0, len(origins), BATCH):
        chunk = origins[i:i + BATCH]
        hist = np.stack([Y[o - SEQ:o] for o in chunk]).transpose(0, 2, 1)  # (B, 27, 672)
        hist = np.nan_to_num(hist, nan=0.0)                   # mask channel carries validity
        hmask = np.stack([M[o - SEQ:o] for o in chunk]).transpose(0, 2, 1)
        xhist = np.stack([X[o - SEQ:o] for o in chunk])       # (B, 672, F)
        xfut = np.stack([X[o + HORIZON:o + 2 * HORIZON] for o in chunk])
        y = np.stack([Y[o + HORIZON:o + 2 * HORIZON] for o in chunk]).transpose(0, 2, 1)  # (B,27,96)
        m = np.stack([M[o + HORIZON:o + 2 * HORIZON] for o in chunk]).transpose(0, 2, 1)
        yield hist, hmask, xhist, xfut, y, m


def build_model(args, n_features, n_consumers, n_nodes):
    import torch.nn as nn
    from heatingforecast.distributions import JohnsonSU

    class Net(nn.Module):
        def __init__(self):
            super().__init__()
            self.encoder = nn.LSTM(2 + n_features, HIDDEN, num_layers=2,
                                   batch_first=True, dropout=0.1)
            self.is_gnn = args.model == "gnn"
            if self.is_gnn:
                self.junction_emb = nn.Parameter(torch.zeros(1, n_nodes - n_consumers, HIDDEN))
                self.gn1 = nn.Linear(HIDDEN, HIDDEN)
                self.gn2 = nn.Linear(HIDDEN, HIDDEN)
                self.norm1 = nn.LayerNorm(HIDDEN)
                self.norm2 = nn.LayerNorm(HIDDEN)
            self.decoder = nn.Sequential(
                nn.Linear(2 * HIDDEN, 256), nn.GELU(), nn.Linear(256, 256), nn.GELU(),
                nn.Linear(256, len(LEVELS) + 4))

        def forward(self, hist, hmask, xhist, xfut):
            B, N, S = hist.shape
            F = xhist.shape[2]
            T = S + HORIZON
            xh = xhist[:, None].expand(B, N, S, F)
            past = torch.cat([hist.unsqueeze(-1), hmask.unsqueeze(-1), xh], dim=-1)
            fut = torch.cat([
                torch.zeros(B, N, HORIZON, 2, device=hist.device),
                xfut[:, None].expand(B, N, HORIZON, F)], dim=-1)
            seq = torch.cat([past, fut], dim=2).reshape(B * N, T, 2 + F)
            enc, _ = self.encoder(seq)
            f = enc[:, -HORIZON:].reshape(B, N, HORIZON, HIDDEN)   # aligned future states
            h = f.mean(dim=2)                                      # node summary
            if self.is_gnn:
                A = self.A
                nodes = torch.cat([h, self.junction_emb.expand(B, -1, -1)], dim=1)
                z = torch.relu(self.gn1(torch.matmul(A, nodes)))
                nodes = self.norm1(nodes + z)
                z = torch.relu(self.gn2(torch.matmul(A, nodes)))
                nodes = self.norm2(nodes + z)
                h = nodes[:, :N]
            h = h[:, :, None, :].expand(B, N, HORIZON, HIDDEN)
            out = self.decoder(torch.cat([h, f], dim=-1))           # (B,N,96,K+4)
            q = out[..., :len(LEVELS)]
            loc = out[..., len(LEVELS)]
            scale = torch.nn.functional.softplus(out[..., len(LEVELS) + 1]) + 0.02
            skew = out[..., len(LEVELS) + 2]
            tail = (torch.nn.functional.softplus(out[..., len(LEVELS) + 3]) + 0.05).clamp(max=8.0)
            return q, JohnsonSU(loc=loc, scale=scale, skew=skew, tail=tail)

    return Net()


def masked_pinball(q, y, m, levels_t):
    # q (B,N,96,K), y (B,N,96), m (B,N,96); NaN targets are masked out
    y = torch.nan_to_num(y, nan=0.0)
    l = levels_t.view(1, 1, 1, -1)
    diff = y.unsqueeze(-1) - q
    loss = torch.maximum(l * diff, (l - 1.0) * diff)               # (B,N,96,K)
    loss = loss.mean(-1) * m
    return loss.sum() / m.sum().clamp(min=1.0)


def masked_nll(dist, y, m):
    y = torch.nan_to_num(y, nan=0.0)
    nll = -dist.log_prob(y)
    # extreme batches (scale collapse / tail growth) must not destabilize training
    nll = torch.clamp(nll, min=-20.0, max=100.0)
    nll = torch.nan_to_num(nll, nan=0.0, posinf=100.0, neginf=-20.0) * m
    return nll.sum() / m.sum().clamp(min=1.0)


def evaluate(D, model, origins, device):
    from heatingforecast.metrics import evaluate_forecast
    model.eval()
    qs, ys, ms, jsu_crps = [], [], [], None
    with torch.no_grad():
        for hist, hmask, xhist, xfut, y, m in make_batches(origins, D):
            tb = torch.tensor(hist, device=device)
            tm = torch.tensor(hmask, device=device)
            tx = torch.tensor(xhist, device=device)
            tf = torch.tensor(xfut, device=device)
            q, dist = model(tb, tm, tx, tf)
            qs.append(q.cpu().numpy().transpose(0, 2, 1, 3).reshape(-1, 27, len(LEVELS)))
            ys.append(y.transpose(0, 2, 1).reshape(-1, 27))
            ms.append(m.transpose(0, 2, 1).reshape(-1, 27))
    q = np.concatenate(qs); y = np.concatenate(ys); m = np.concatenate(ms)
    mask = (m > 0) & ~np.isnan(y)
    res = evaluate_forecast(y[mask], q[mask], LEVELS)
    per_c = {f"CONS_{c+1:03d}": round(evaluate_forecast(
        y[mask[:, c], c], q[mask[:, c], c], LEVELS)["pinball"], 4)
        for c in range(27) if mask[:, c].sum() > 100}
    res["pinball_per_consumer"] = per_c
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", choices=["lstm", "gnn"],
                    default=os.environ.get("SEQ_MODEL", "gnn"))
    ap.add_argument("--smoke", action="store_true",
                    default=os.environ.get("SEQ_SMOKE", "") == "1")
    args = ap.parse_args()

    torch.manual_seed(SEED)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"device: {device}", flush=True)

    root = Path.cwd()
    files = download_all(root)
    D = prepare(files)
    import pandas as pd
    index = pd.date_range("2024-12-31 23:00", periods=len(D["Y"]), freq="15min", tz="UTC")
    origins_fit = daily_origins(index, D["fit"])
    origins_test = daily_origins(index, D["test"])
    n_val = N_VAL_ORIGINS if not args.smoke else 4
    origins_tr, origins_va = origins_fit[:-n_val], origins_fit[-n_val:]
    if args.smoke:
        origins_tr, origins_test = origins_tr[::30], origins_test[:2]
    print(f"origins: train {len(origins_tr)}, val {len(origins_va)}, test {len(origins_test)}",
          flush=True)

    n_features = D["X"].shape[1]
    n_consumers = D["Y"].shape[1]
    model = build_model(args, n_features, n_consumers, D["A"].shape[0]).to(device)
    model.A = torch.tensor(D["A"], device=device) if args.model == "gnn" else None
    n_par = sum(p.numel() for p in model.parameters())
    print(f"model: {args.model}, {n_par} params", flush=True)

    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(
        opt, T_max=(5 if args.smoke else EPOCHS))
    levels_t = torch.tensor(LEVELS, device=device)

    best_val, best_state, t0 = float("inf"), None, time.time()
    epochs = 3 if args.smoke else EPOCHS
    for ep in range(epochs):
        model.train()
        te, tn, nb = 0.0, 0.0, 0
        for hist, hmask, xhist, xfut, y, m in make_batches(origins_tr, D):
            tb, tm_ = torch.tensor(hist, device=device), torch.tensor(hmask, device=device)
            tx = torch.tensor(xhist, device=device)
            tf = torch.tensor(xfut, device=device)
            ty, tmy = torch.tensor(y, device=device), torch.tensor(m, device=device)
            q, dist = model(tb, tm_, tx, tf)
            loss = masked_pinball(q, ty, tmy, levels_t) + 0.5 * masked_nll(dist, ty, tmy)
            if not torch.isfinite(loss):
                mq = tmy > 0
                print(f"  [skip] non-finite loss: max|y|={ty[mq].abs().max().item():.1f} "
                      f"max|q|={q.abs().max().item():.2e} "
                      f"scale=[{dist.scale.min().item():.2e},{dist.scale.max().item():.2e}] "
                      f"tail=[{dist.tail.min().item():.2f},{dist.tail.max().item():.2f}]",
                      flush=True)
                opt.zero_grad()
                continue
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            if any(p.grad is not None and not torch.isfinite(p.grad).all()
                   for p in model.parameters()):
                print("  [skip] non-finite gradients — step suppressed", flush=True)
                opt.zero_grad()
                continue
            opt.step()
            te += float(masked_pinball(q, ty, tmy, levels_t).detach())
            tn += float(loss.detach()); nb += 1
        sched.step()
        # validation
        model.eval()
        with torch.no_grad():
            vloss, vn = 0.0, 0.0
            for hist, hmask, xhist, xfut, y, m in make_batches(origins_va, D):
                tb, tm_ = torch.tensor(hist, device=device), torch.tensor(hmask, device=device)
                tx = torch.tensor(xhist, device=device)
                tf = torch.tensor(xfut, device=device)
                ty, tmy = torch.tensor(y, device=device), torch.tensor(m, device=device)
                q, dist = model(tb, tm_, tx, tf)
                vl = masked_pinball(q, ty, tmy, levels_t).item()
                vloss += vl; vn += 1
        vloss /= max(vn, 1)
        marker = ""
        if vloss < best_val:
            best_val, best_state = vloss, copy.deepcopy(model.state_dict())
            marker = " *"
        print(f"epoch {ep+1}: train pinball {te/max(nb,1):.4f} loss {tn/max(nb,1):.4f} "
              f"val pinball {vloss:.4f}{marker} ({time.time()-t0:.0f}s)", flush=True)
        if os.environ.get("TRACKIO_SPACE_ID"):
            import trackio
            trackio.log({"train_pinball": te / max(nb, 1), "val_pinball": vloss},
                        step=ep + 1)

    if best_state is None:  # every val was NaN — fall back to final weights
        best_state = copy.deepcopy(model.state_dict())
    model.load_state_dict(best_state)
    res = evaluate(D, model, origins_test, device)
    print(f"TEST {args.model}: {json.dumps({k: v for k, v in res.items() if not isinstance(v, dict)}, indent=None)}", flush=True)

    if args.smoke:
        return
    out = root / "results"
    out.mkdir(exist_ok=True)
    (out / f"{args.model}.json").write_text(json.dumps(
        {"model": args.model, "params": n_par, "epochs": epochs,
         "best_val_pinball_scaled": best_val,
         "protocol": "day-ahead, origin midnight UTC, 7-day history, "
                     "weather perfect-foresight",
         "test": res}, indent=2))
    (root / "models").mkdir(exist_ok=True)
    torch.save({"state_dict": best_state, "model": args.model,
                "hidden": HIDDEN, "n_features": n_features,
                "mu": D["mu"], "sd": D["sd"], "A": D["A"],
                "consumer_nodes": D["consumer_nodes"]},
               root / "models" / f"{args.model}.pt")
    from huggingface_hub import HfApi
    api = HfApi()
    for sub in ["results", "models"]:
        api.upload_file(path_or_fileobj=root / sub / f"{args.model}{'.'+'json' if sub=='results' else '.pt'}",
                        path_in_repo=f"{sub}/{args.model}.{'json' if sub=='results' else 'pt'}",
                        repo_id=REPO_CODE, repo_type="model",
                        commit_message=f"{args.model} results + weights")
    print("pushed results + weights", flush=True)


if __name__ == "__main__":
    main()