import json
import random
import sys
import numpy as np
import torch
from torch import nn

sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
from bench import make_model, train_model, evaluate, sweep_baseline, make_report
from manifold_dynamics_track import get_dataset

SEEDS = tuple(range(8))
NTR, NTE = 400, 200
EPOCHS = 18
BATCH = 128
LRS = [1e-3, 3e-3, 1e-2]


def seed_all(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def dataset(seed):
    d = get_dataset(seed, NTR, NTE)
    for k in ("xtr", "ytr", "xte", "yte"):
        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
    d["input_shape"] = tuple(d["xtr"].shape[1:])
    d["out_dim"] = 3
    return d


def baseline_factory(cfg):
    def run(seed):
        seed_all(seed)
        d = dataset(seed)
        net = make_model("rnn_small", d["input_shape"], 3)
        _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"], batch=BATCH,
                                   log=lambda *_: None)
        return float(metric)
    return run


def intrinsic_factory(cfg, collect=False):
    def run(seed):
        seed_all(seed)
        d = dataset(seed)
        net = make_model("rnn_small", d["input_shape"], 3)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        try:
            net = net.to(device)
            xtr, ytr = d["xtr"].to(device), d["ytr"].to(device)
            xte, yte = d["xte"].to(device), d["yte"].to(device)
            opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"])
            for _ in range(EPOCHS):
                net.train()
                perm = torch.randperm(len(xtr), device=device)
                for j in range(0, len(xtr), BATCH):
                    ix = perm[j:j+BATCH]
                    raw = net(xtr[ix])
                    # Intrinsic S1 controller/state: retract the angular pair.
                    emb = raw[:, :2]
                    emb = emb / torch.clamp(torch.linalg.vector_norm(emb, dim=1, keepdim=True), min=1e-7)
                    pred = torch.cat((emb, raw[:, 2:3]), dim=1)
                    loss = ((pred - ytr[ix]) ** 2).mean()
                    opt.zero_grad(set_to_none=True)
                    loss.backward()
                    opt.step()
            net.eval()
            with torch.no_grad():
                raw = net(xte)
                emb = raw[:, :2] / torch.clamp(torch.linalg.vector_norm(raw[:, :2], dim=1, keepdim=True), min=1e-7)
                pred = torch.cat((emb, raw[:, 2:3]), dim=1)
                metric = float(((pred - yte) ** 2).mean())
                violation = float(torch.abs(torch.linalg.vector_norm(pred[:, :2], dim=1) - 1).max())
            if collect:
                return metric, net, d, violation
            return metric
        except RuntimeError:
            # Explicit CPU fallback, matching the bench's robust device policy.
            seed_all(seed)
            net = make_model("rnn_small", d["input_shape"], 3).to("cpu")
            xtr, ytr = d["xtr"], d["ytr"]
            opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"])
            for _ in range(EPOCHS):
                perm = torch.randperm(len(xtr))
                for j in range(0, len(xtr), BATCH):
                    ix = perm[j:j+BATCH]
                    raw = net(xtr[ix]); emb = raw[:, :2] / torch.clamp(torch.linalg.vector_norm(raw[:, :2], dim=1, keepdim=True), min=1e-7)
                    pred = torch.cat((emb, raw[:, 2:3]), 1)
                    loss = ((pred-ytr[ix])**2).mean()
                    opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
            with torch.no_grad():
                raw = net(d["xte"]); emb = raw[:, :2] / torch.clamp(torch.linalg.vector_norm(raw[:, :2], dim=1, keepdim=True), min=1e-7)
                pred = torch.cat((emb, raw[:, 2:3]), 1)
                metric = float(((pred-d["yte"])**2).mean())
                violation = float(torch.abs(torch.linalg.vector_norm(pred[:, :2], dim=1)-1).max())
            return (metric, net, d, violation) if collect else metric
    return run


def signature(cfg):
    # Evaluate both trained systems on the same held-out examples.
    s = 0
    seed_all(s); d = dataset(s)
    bnet, _, _ = train_model(make_model("rnn_small", d["input_shape"], 3), d,
                             epochs=EPOCHS, lr=cfg["lr"], batch=BATCH, log=lambda *_: None)
    bdev = next(bnet.parameters()).device
    with torch.no_grad():
        raw = bnet(d["xte"].to(bdev))
        bviol = float(torch.abs(torch.linalg.vector_norm(raw[:, :2], dim=1)-1).max())
    imetric, _, _, iviol = intrinsic_factory(cfg, collect=True)(s)
    return {
        "prediction": "intrinsic retraction keeps every predicted embedded angular state on S1; Euclidean output has nonzero norm error",
        "predicted_baseline_violation_order": "nonzero",
        "predicted_idea_violation": 0.0,
        "observed_baseline_max_violation": bviol,
        "observed_idea_max_violation": iviol,
        "observed_idea_metric_seed0": imetric,
        "confirmed": bool(bviol > 1e-6 and iviol < 1e-5)
    }


def main():
    grid = [{"lr": x} for x in LRS]
    base = sweep_baseline(baseline_factory, grid, seeds=(0, 1, 2, 3))
    trials = [{"cfg": c, "result": evaluate(intrinsic_factory(c), SEEDS)} for c in grid]
    best = min(trials, key=lambda z: z["result"]["mean"])
    extra = {
        "custom_track": {"name": "manifold_pendulum", "file": "manifold_dynamics_track.py", "domain": "dynamics_and_embedded_manifolds"},
        "idea_config": best["cfg"],
        "idea_sweep": trials,
        "mechanism_signature": signature(best["cfg"])
    }
    rep = make_report("manifold_pendulum", "rnn_small", base, best["result"], extra)
    with open("bench_report.json", "w") as f:
        json.dump(rep, f, indent=2)
    print(json.dumps(rep, indent=2))

if __name__ == "__main__":
    main()