import sys, json, math, random
import numpy as np
import torch
import torch.nn as nn

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

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
EPOCHS = 15
NTRAIN, NTEST = 800, 300
BATCH = 128

# Same GRU backbone and hidden size as bench.models.rnn_small.
class MixtureRNN(nn.Module):
    def __init__(self, hidden=64, modes=2):
        super().__init__()
        self.rnn = nn.GRU(3, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 3*modes)  # logits, means, log standard deviations
        self.modes = modes
    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        _, h = self.rnn(seq)
        return self.head(h[-1])

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        try: torch.cuda.manual_seed_all(seed)
        except Exception: pass

def device():
    return "cuda" if torch.cuda.is_available() else "cpu"

def mixture_loss(raw, y):
    k = raw.shape[1] // 3
    logits, means, logstd = raw[:, :k], raw[:, k:2*k], raw[:, 2*k:]
    logstd = logstd.clamp(-5.0, 2.0)
    z = (y - means) / logstd.exp()
    lp = -0.5*z*z - logstd - 0.5*math.log(2*math.pi)
    return -(torch.log_softmax(logits, 1) + lp).logsumexp(1).mean()

def train_mixture(seed, lr, epochs=EPOCHS, return_model=False):
    seed_all(seed); ds = get_dataset("dynamics", seed, NTRAIN, NTEST)
    net = MixtureRNN(); dev = device()
    try:
        net.to(dev); x, y = ds["xtr"].to(dev), ds["ytr"].to(dev)
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        for _ in range(epochs):
            net.train(); perm = torch.randperm(len(x), device=dev)
            for i in range(0, len(x), BATCH):
                ix = perm[i:i+BATCH]; loss = mixture_loss(net(x[ix]), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad():
            raw = net(ds["xte"].to(dev)); k=2
            w = torch.softmax(raw[:,:k],1); mu=raw[:,k:2*k]
            pred=(w*mu).sum(1,keepdim=True)
            mse=((pred-ds["yte"].to(dev))**2).mean().item()
        return (mse, net, ds, raw.detach().cpu()) if return_model else mse
    except RuntimeError:
        # Small CPU retry mirrors the harness's GPU-fallback intent.
        seed_all(seed); net = MixtureRNN(); net.to("cpu")
        x,y=ds["xtr"],ds["ytr"]; opt=torch.optim.Adam(net.parameters(),lr=lr)
        for _ in range(epochs):
            perm=torch.randperm(len(x))
            for i in range(0,len(x),BATCH):
                ix=perm[i:i+BATCH]; loss=mixture_loss(net(x[ix]),y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            raw=net(ds["xte"]); w=torch.softmax(raw[:,:2],1); mu=raw[:,2:4]
            mse=((w*mu).sum(1,keepdim=True)-ds["yte"]).pow(2).mean().item()
        return (mse,net,ds,raw) if return_model else mse

def baseline_fn(cfg):
    def run(seed):
        seed_all(seed); ds=get_dataset("dynamics",seed,NTRAIN,NTEST)
        net=make_model("rnn_small",ds["input_shape"],ds["out_dim"])
        _, metric, _=train_model(net,ds,epochs=EPOCHS,lr=cfg["lr"],batch=BATCH,log=lambda *_:None)
        return metric
    return run

def idea_fn(cfg):
    return lambda seed: train_mixture(seed,cfg["lr"])

def signature(seed, lr):
    mse, net, ds, raw = train_mixture(seed,lr,return_model=True)
    w=torch.softmax(raw[:,:2],1); mu=raw[:,2:4]; sd=raw[:,4:6].clamp(-5,2).exp()
    sep=(mu[:,0]-mu[:,1]).abs(); avg_sd=(w*sd).sum(1)
    # Model-derived chance estimate at theta<=0 versus observed test frequency.
    cdf=0.5*(1+torch.erf((-mu)/(sd*math.sqrt(2))))
    p_mix=(w*cdf).sum(1).numpy(); p_gauss=(0.5*(1+torch.erf((-(w*mu).sum(1))/(torch.sqrt((w*(sd**2+(mu-(w*mu).sum(1,keepdim=True))**2)).sum(1))*math.sqrt(2))))).numpy()
    obs=(ds["yte"].numpy().reshape(-1)<=0).astype(float)
    return {"n":len(obs),"test_mse":mse,"predicted_mean_mode_separation":float(sep.mean()),"predicted_mean_component_sd":float(avg_sd.mean()),"predicted_bimodal_fraction":float((sep>2*avg_sd).float().mean()),"mixture_chance_abs_error":float(abs(p_mix.mean()-obs.mean())),"moment_gaussian_chance_abs_error":float(abs(p_gauss.mean()-obs.mean())),"observed_event_rate":float(obs.mean()),"confirmed":bool(sep.mean().item()>0.05 and abs(p_mix.mean()-obs.mean()) < abs(p_gauss.mean()-obs.mean()))}

def main():
    grid=[{"lr":1e-3},{"lr":3e-3},{"lr":1e-2}]
    base=sweep_baseline(baseline_fn,grid,seeds=SWEEP_SEEDS)
    # Equal-size intervention sweep over exactly the baseline union of learning rates.
    idea_trials=[]
    for cfg in grid:
        r=evaluate(idea_fn(cfg),seeds=SWEEP_SEEDS)
        idea_trials.append({"cfg":cfg,"mean":r["mean"]})
    best_cfg=min(idea_trials,key=lambda z:z["mean"])["cfg"]
    idea=evaluate(idea_fn(best_cfg),seeds=SEEDS)
    # Nearby settings are explicitly run on all paired seeds for transparent reporting.
    nearby={str(c["lr"]):evaluate(idea_fn(c),seeds=SEEDS) for c in grid}
    base["idea_union_sweep_note"]="baseline evaluated at every intervention learning rate; final baseline is its tuned best config"
    rep=make_report("dynamics","rnn_small",base,idea,extra={"prediction":"A learned mixture should retain separated predictive modes and improve threshold-probability calibration when the trained task is multimodal.","idea_sweep":idea_trials,"idea_nearby_full":nearby,"trained_model_signature":signature(SEEDS[0],best_cfg["lr"])})
    rep["selection"]={"idea_best_cfg":best_cfg,"epochs":EPOCHS,"n_train":NTRAIN,"n_test":NTEST,"structural_match":"dynamics: controlled pendulum multi-step target"}
    with open("bench_report.json","w") as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))

if __name__=="__main__": main()