"""Trains a convolutional network that predicts the sheet parameters from the
spectrogram of a strike, evaluates it on held-out clips, and writes a report.

    python3 ml/train.py --data DATASET_DIR --out RUN_DIR [--cache CACHE_DIR] [--epochs N] [--val-every K] [--device cuda|cpu]

Targets are the logarithms of E, rho, h, Lx, Ly, standardized by the training
set; the loss is the mean squared error over the five. Every K-th clip by
index is held out for validation. The report gives, per target and for a few
derived quantities (E/rho, the bending wave constant kappa, the area, the
aspect ratio), the relative RMS error of the prediction against that of
predicting the training mean, and scatter plots. A line is appended to
RUN_DIR/../history.jsonl so that the error can be tracked as clips accumulate.
"""

import argparse
import base64
import io
import json
import math
import os
import sys
import time

import numpy as np
import torch
import torch.nn as nn

# The office machine's PyTorch build cannot find a cuDNN engine for the
# convolution backward pass; the native kernels work, more slowly.
torch.backends.cudnn.enabled = os.environ.get("CUDNN", "0") == "1"

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from features import build_cache  # noqa: E402

TARGETS = ["E", "rho", "h", "Lx", "Ly"]
UNITS = {"E": ("GPa", 1e-9), "rho": ("kg/m^3", 1), "h": ("mm", 1e3), "Lx": ("in", 1 / 0.0254), "Ly": ("in", 1 / 0.0254)}
DERIVED = {
    "E/rho": lambda d: d["E"] - d["rho"],
    "kappa": lambda d: d["h"] + 0.5 * (d["E"] - d["rho"]),
    "area": lambda d: d["Lx"] + d["Ly"],
    "aspect": lambda d: d["Lx"] - d["Ly"],
    "kappa/area": lambda d: d["h"] + 0.5 * (d["E"] - d["rho"]) - d["Lx"] - d["Ly"],
}


class Net(nn.Module):
    """Six strided conv blocks, global average pooling, and a small head. A
    second input channel carries the frequency coordinate of each row, so the
    network can read absolute frequency despite the pooling."""

    def __init__(self, n_out=5, channels=(16, 32, 64, 96, 128, 128)):
        super().__init__()
        layers, c_in = [], 2
        for c in channels:
            layers += [nn.Conv2d(c_in, c, 3, padding=1), nn.BatchNorm2d(c), nn.ReLU(inplace=True),
                       nn.Conv2d(c, c, 3, padding=1), nn.BatchNorm2d(c), nn.ReLU(inplace=True), nn.MaxPool2d(2)]
            c_in = c
        self.features = nn.Sequential(*layers)
        self.head = nn.Sequential(nn.Linear(c_in, 128), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(128, n_out))

    def forward(self, x):
        b, _, f, t = x.shape
        coord = torch.linspace(0, 1, f, device=x.device).view(1, 1, f, 1).expand(b, 1, f, t)
        h = self.features(torch.cat([x, coord], 1)).mean((2, 3))
        return self.head(h)


class SpecNet(nn.Module):
    """Full-frequency-resolution model. The spectrogram is collapsed into a few
    time slices (mean log spectrum over 0-0.25, 0.25-1, 1-2, 2-3 s), which keep
    the partials at the resolution of the bins and the decay of each band as
    the ratio between slices. A 1D convolutional stack along frequency (two
    stride-2 steps only, dilations for context) is followed by multi-head
    attention pooling over frequency positions with a sinusoidal position code,
    so that the pooled vector can express where the lines are."""

    def __init__(self, n_out=5, n_freq=1024, fps=48000 / 1024, slices=((0, 0.25), (0.25, 1), (1, 2), (2, 3)), heads=4, pos_dim=16):
        super().__init__()
        self.bounds = [(int(a * fps), int(b * fps)) for a, b in slices]
        c_in = len(slices) + 1
        def block(ci, co, k, stride=1, dil=1):
            return [nn.Conv1d(ci, co, k, stride=stride, padding=dil * (k // 2), dilation=dil), nn.BatchNorm1d(co), nn.ReLU(inplace=True)]
        self.conv = nn.Sequential(*block(c_in, 32, 7), *block(32, 32, 7), *block(32, 64, 5, stride=2), *block(64, 64, 5, dil=2),
                                  *block(64, 96, 5, stride=2), *block(96, 96, 5, dil=4), *block(96, 96, 5, dil=8), *block(96, 128, 3))
        n_pos = (n_freq + 3) // 4
        pos = torch.arange(n_pos, dtype=torch.float32)[:, None] / n_pos
        freqs = 2 ** torch.arange(pos_dim // 2, dtype=torch.float32)
        self.register_buffer("pos", torch.cat([torch.sin(math.pi * pos * freqs), torch.cos(math.pi * pos * freqs)], 1))
        d = 128 + pos_dim
        self.att = nn.Linear(d, heads)
        self.head = nn.Sequential(nn.Linear(heads * d, 256), nn.ReLU(inplace=True), nn.Dropout(0.2), nn.Linear(256, n_out))

    def forward(self, x):
        b, _, f, t = x.shape
        s = torch.stack([x[:, 0, :, a:min(b_, t)].mean(-1) for a, b_ in self.bounds], 1)
        coord = torch.linspace(0, 1, f, device=x.device).view(1, 1, f).expand(b, 1, f)
        h = self.conv(torch.cat([s, coord], 1)).transpose(1, 2)  # (b, positions, 128)
        hp = torch.cat([h, self.pos[: h.shape[1]].expand(b, -1, -1)], -1)
        a = torch.softmax(self.att(hp), dim=1)  # (b, positions, heads)
        pooled = torch.einsum("bpk,bpd->bkd", a, hp).flatten(1)
        return self.head(pooled)


# Alternative target basis: the combinations the linear physics identifies
# (kappa / area up to a constant, aspect, E / rho, area) plus rho, so that the
# well-determined directions get their own outputs and clean gradients.
COMBOS = np.array([
    [0.5, -0.5, 1, -1, -1],  # log kappa/area (up to a constant)
    [0, 0, 0, 1, -1],  # log aspect
    [1, -1, 0, 0, 0],  # log E/rho
    [0, 0, 0, 1, 1],  # log area
    [0, 1, 0, 0, 0],  # log rho
], dtype=np.float32)


def metrics(pred, true, mean):
    """Relative RMS error (percent) of pred and of the constant mean, and R^2, all in log space."""
    err = np.sqrt(np.mean((pred - true) ** 2))
    base = np.sqrt(np.mean((mean - true) ** 2))
    return {"rel_rms_pct": float(100 * err), "baseline_rel_rms_pct": float(100 * base), "r2": float(1 - err ** 2 / max(base ** 2, 1e-12)), "mae_pct": float(100 * np.mean(np.abs(pred - true)))}


def scatter_png(true, pred, name, unit, scale):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(3.2, 3.2), dpi=100)
    t, p = np.exp(true) * scale, np.exp(pred) * scale
    ax.scatter(t, p, s=6, alpha=0.5)
    lo, hi = min(t.min(), p.min()), max(t.max(), p.max())
    ax.plot([lo, hi], [lo, hi], "k-", lw=0.8)
    ax.set_xlabel(f"true {name} [{unit}]")
    ax.set_ylabel(f"predicted {name} [{unit}]")
    ax.set_xscale("log"); ax.set_yscale("log")
    fig.tight_layout()
    buf = io.BytesIO(); fig.savefig(buf, format="png"); plt.close(fig)
    return base64.b64encode(buf.getvalue()).decode()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--cache", default=None)
    ap.add_argument("--epochs", type=int, default=60)
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--val-every", type=int, default=8, help="every k-th clip by index is held out")
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--logf", action="store_true", help="log-frequency spectrogram input")
    ap.add_argument("--no-augment", action="store_true", help="disable the onset-shift and masking augmentation")
    ap.add_argument("--arch", default="cnn2d", choices=["cnn2d", "spec1d"], help="pooling 2D CNN (default) or the full-resolution SpecNet")
    ap.add_argument("--targets", default="params", choices=["params", "combos"], help="train on log E, rho, h, Lx, Ly or on the identifiable combinations")
    ap.add_argument("--tag", default=None, help="suffix of the report and history files (default: 'logf' with --logf, else none)")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    kind = "logf" if args.logf else "lin"
    tag = args.tag if args.tag is not None else "-".join(t for t in [("logf" if args.logf else ""), ("" if args.arch == "cnn2d" else args.arch), ("" if args.targets == "params" else args.targets)] if t)
    suffix = f"-{tag}" if tag else ""
    cache = args.cache or os.path.join(args.data, "ml", "cache" + ("-logf" if args.logf else ""))
    t_start = time.time()
    rows, feats = build_cache(args.data, cache, args.workers, kind)
    n = len(rows)
    y = np.stack([[math.log(r[t]) for t in TARGETS] for r in rows]).astype(np.float32)
    idx = np.array([r["index"] for r in rows])
    val = idx % args.val_every == 0
    tr = ~val
    if tr.sum() < 20 or val.sum() < 5:
        print(f"only {n} clips; need more"); return
    T = COMBOS if args.targets == "combos" else np.eye(len(TARGETS), dtype=np.float32)
    Tinv = np.linalg.inv(T)
    yt = y @ T.T
    mu, sd = yt[tr].mean(0), yt[tr].std(0) + 1e-6
    z = (yt - mu) / sd
    print(f"{n} clips: {tr.sum()} train, {val.sum()} val; features {feats.shape}; device {args.device}")

    dev = torch.device(args.device)
    X = torch.from_numpy(feats)  # float16 (n, F, T), kept on the CPU
    Z = torch.from_numpy(z)
    net = (SpecNet(len(TARGETS), n_freq=feats.shape[1]) if args.arch == "spec1d" else Net(len(TARGETS))).to(dev)
    print(f"arch {args.arch}, targets {args.targets}, {sum(p.numel() for p in net.parameters()) / 1e6:.2f} M parameters")
    opt = torch.optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=args.lr, total_steps=args.epochs * math.ceil(tr.sum() / args.batch))
    tr_idx, val_idx = np.where(tr)[0], np.where(val)[0]
    rng = np.random.default_rng(0)

    def predict(ids):
        net.eval()
        out = []
        with torch.no_grad():
            for k in range(0, len(ids), 64):
                xb = X[ids[k:k + 64]].to(dev).float().unsqueeze(1)
                out.append(net(xb).cpu().numpy())
        return np.concatenate(out)

    def augment(xb):
        """Label-preserving perturbations against memorization: a random onset
        shift of up to 8 frames (170 ms), and one masked frequency band and one
        masked time span per clip (the masked cells take the floor value 0)."""
        bsz, _, nf, nt = xb.shape
        out = torch.zeros_like(xb)
        for i in range(bsz):
            sh = int(rng.integers(0, 9))
            out[i, :, :, sh:] = xb[i, :, :, : nt - sh]
            f0, fw = int(rng.integers(0, nf)), int(rng.integers(0, nf // 16))
            out[i, :, f0:f0 + fw, :] = 0
            t0, tw = int(rng.integers(0, nt)), int(rng.integers(0, nt // 8))
            out[i, :, :, t0:t0 + tw] = 0
        return out

    history, best, best_state = [], float("inf"), None
    for epoch in range(args.epochs):
        net.train()
        perm = rng.permutation(tr_idx)
        tot = 0.0
        for k in range(0, len(perm), args.batch):
            b = perm[k:k + args.batch]
            xb = X[b].to(dev).float().unsqueeze(1)
            if not args.no_augment:
                xb = augment(xb)
            zb = Z[b].to(dev)
            loss = nn.functional.mse_loss(net(xb), zb)
            opt.zero_grad(); loss.backward(); opt.step(); sched.step()
            tot += loss.item() * len(b)
        pv = predict(val_idx)
        vloss = float(np.mean((pv - z[val_idx]) ** 2))
        history.append({"epoch": epoch, "train_loss": tot / len(perm), "val_loss": vloss})
        if vloss < best:
            best, best_state = vloss, {k: v.detach().cpu().clone() for k, v in net.state_dict().items()}
        print(f"epoch {epoch}: train {tot / len(perm):.4f} val {vloss:.4f}", flush=True)
    net.load_state_dict(best_state)
    torch.save({"state": best_state, "mu": mu, "sd": sd, "targets": TARGETS, "T": T, "arch": args.arch, "input": kind, "n_freq": feats.shape[1]}, os.path.join(args.out, "model.pt"))

    # Evaluation in log space, back in the natural scale of each target.
    pv = (predict(val_idx) * sd + mu) @ Tinv.T
    yv = y[val_idx]
    ym = np.broadcast_to(mu @ Tinv.T, yv.shape)
    res = {t: metrics(pv[:, i], yv[:, i], ym[:, i]) for i, t in enumerate(TARGETS)}
    d_pred = {t: pv[:, i] for i, t in enumerate(TARGETS)}
    d_true = {t: yv[:, i] for i, t in enumerate(TARGETS)}
    d_mean = {t: ym[:, i] for i, t in enumerate(TARGETS)}
    for name, f in DERIVED.items():
        res[name] = metrics(f(d_pred), f(d_true), f(d_mean))
    summary = {"input": kind, "arch": args.arch, "targets": args.targets, "augment": not args.no_augment, "n_clips": n, "n_train": int(tr.sum()), "n_val": int(val.sum()), "best_val_loss": float(best), "epochs": args.epochs,
               "minutes": (time.time() - t_start) / 60, "finished": time.strftime("%Y-%m-%d %H:%M:%S"), "metrics": res, "history": history}
    with open(os.path.join(args.out, "metrics.json"), "w") as f:
        json.dump(summary, f, indent=1)
    with open(os.path.join(args.out, "..", f"history{suffix}.jsonl"), "a") as f:
        f.write(json.dumps({k: v for k, v in summary.items() if k != "history"}) + "\n")
    np.save(os.path.join(args.out, "val_pred.npy"), np.stack([idx[val_idx], *yv.T, *pv.T], 1))
    write_report(args.out, summary, yv, pv, suffix)
    print("\n".join(f"{t:>8}: rel RMS {m['rel_rms_pct']:.1f}% (mean-baseline {m['baseline_rel_rms_pct']:.1f}%), R^2 {m['r2']:.2f}" for t, m in res.items()))


def write_report(out, summary, yv, pv, suffix=""):
    hist_lines = []
    hp = os.path.join(out, "..", f"history{suffix}.jsonl")
    if os.path.exists(hp):
        with open(hp) as f:
            hist_lines = [json.loads(l) for l in f if l.strip()]
    rows = "".join(
        f"<tr><td>{t}</td><td>{m['rel_rms_pct']:.1f}</td><td>{m['mae_pct']:.1f}</td><td>{m['baseline_rel_rms_pct']:.1f}</td><td>{m['r2']:.2f}</td></tr>"
        for t, m in summary["metrics"].items())
    imgs = "".join(f'<img src="data:image/png;base64,{scatter_png(yv[:, i], pv[:, i], t, *UNITS[t])}">' for i, t in enumerate(TARGETS))
    hist_rows = "".join(f"<tr><td>{h['finished']}</td><td>{h['n_clips']}</td>" + "".join(f"<td>{h['metrics'][t]['rel_rms_pct']:.1f}</td>" for t in TARGETS) + "</tr>" for h in hist_lines)
    curve = "".join(f"<tr><td>{h['epoch']}</td><td>{h['train_loss']:.4f}</td><td>{h['val_loss']:.4f}</td></tr>" for h in summary["history"][::max(1, len(summary["history"]) // 15)])
    html = f"""<!doctype html><meta charset="utf-8"><title>steelsim31 prediction report</title>
<style>body{{font:14px system-ui,sans-serif;margin:16px;color:#222}} table{{border-collapse:collapse;margin:8px 0}} td,th{{padding:3px 10px;border-bottom:1px solid #ddd;text-align:right}} td:first-child,th:first-child{{text-align:left}} img{{margin:4px}} h2{{font-size:16px;margin-top:20px}}</style>
<h1 style="font-size:18px">Prediction of the sheet parameters from the strike sound</h1>
<p>Input: {"log-frequency" if summary["input"] == "logf" else "linear-frequency"} spectrogram; architecture {summary.get("arch", "cnn2d")}; targets {summary.get("targets", "params")}. {summary['n_clips']} clips ({summary['n_train']} train, {summary['n_val']} held out), {summary['epochs']} epochs, {summary['minutes']:.1f} min, finished {summary['finished']}. Errors are relative (percent, from the logarithm) on the held-out clips; the baseline predicts the training mean.</p>
<table><tr><th>target</th><th>rel RMS %</th><th>rel MAE %</th><th>baseline rel RMS %</th><th>R²</th></tr>{rows}</table>
<p>Derived: E/rho (sets the wave speeds), kappa = h sqrt(E / (12 rho (1 - nu²))) (sets the bending frequencies), area = Lx Ly, aspect = Lx / Ly, kappa/area (the scale of the bending spectrum).</p>
<h2>Held-out predictions</h2>{imgs}
<h2>Error against the number of clips</h2>
<table><tr><th>run</th><th>clips</th>{"".join(f"<th>{t} %</th>" for t in TARGETS)}</tr>{hist_rows}</table>
<h2>Training curve (standardized MSE)</h2>
<table><tr><th>epoch</th><th>train</th><th>val</th></tr>{curve}</table>
"""
    with open(os.path.join(out, "report.html"), "w") as f:
        f.write(html)
    with open(os.path.join(out, "..", f"report{suffix}.html"), "w") as f:
        f.write(html)


if __name__ == "__main__":
    main()