import sys, json, 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))
EPOCHS = 15
BATCH = 128
# Union of learning rates is shared by baseline and idea-side candidates.
LR_GRID = [1e-3, 3e-3, 1e-2]
AMP_GRID = [0.05, 0.15, 0.30]
GAMMA = 0.15

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 nominal_delta(x, amp, carrier):
    """First-order known pendulum response for the carrier on the final state.
    This is used only as the training consistency target, not as the benchmark metric."""
    z = x.view(-1, 8, 3)
    th, om, u = z[:, -1, 0], z[:, -1, 1], z[:, -1, 2]
    dt = 0.05
    # Four substeps, nominal g and damping, matching the data generator form.
    th0, om0 = th, om
    carr = carrier if torch.is_tensor(carrier) else torch.full_like(u, float(carrier))
    # The benchmark target is a one-step local response at the final observed state.
    # Use the final element of the injected sequence, broadcast over the batch.
    uc = u + carr[:, -1] if carr.ndim == 2 else u + carr
    for _ in range(4):
        om = om + (-9.81 / 10 * torch.sin(th) - .25 * om + 2.0 * uc) * dt / 4
        th = th + om * dt / 4
    for _ in range(4):
        om0 = om0 + (-9.81 / 10 * torch.sin(th0) - .25 * om0 + 2.0 * u) * dt / 4
        th0 = th0 + om0 * dt / 4
    return th - th0

def carrier_tensor(n, amp, device):
    # Structured low-frequency carrier across the 8-step control channel.
    t = torch.arange(8, device=device, dtype=torch.float32)
    c = torch.cos(2.0 * t / 8.0 * 2.0 * np.pi)
    return (amp * c).view(1, 8, 1).expand(n, -1, -1)

def train_carrier(ds, epochs, lr, amp, seed):
    seed_all(seed)
    model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    # Robust explicit device fallback, as this loop is the intervention itself.
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        model = model.to(device)
        opt = torch.optim.Adam(model.parameters(), lr=lr)
        lossf = nn.MSELoss()
        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        for _ in range(epochs):
            model.train(); perm = torch.randperm(len(x), device=device)
            for i in range(0, len(x), BATCH):
                ix = perm[i:i+BATCH]; xb, yb = x[ix], y[ix]
                c = carrier_tensor(len(ix), amp, device)
                xp = xb.view(-1, 8, 3).clone()
                xp[:, :, 2:3] = xp[:, :, 2:3] + c
                # Carrier rollout prediction and passive prediction share weights.
                pred_p = model(xb)
                pred_c = model(xp.reshape(len(ix), -1))
                target_delta = nominal_delta(xb, amp, c[:, :, 0])
                loss = lossf(pred_p, yb) + GAMMA * lossf(pred_c - pred_p, target_delta[:, None])
                opt.zero_grad(); loss.backward(); opt.step()
        model.eval()
        with torch.no_grad():
            out = model(ds['xte'].to(device))
            metric = float(((out - ds['yte'].to(device)) ** 2).mean())
        return metric, model, device
    except RuntimeError:
        model = model.to('cpu')
        opt = torch.optim.Adam(model.parameters(), lr=lr)
        x, y = ds['xtr'], ds['ytr']
        for _ in range(epochs):
            perm = torch.randperm(len(x))
            for i in range(0, len(x), BATCH):
                ix=perm[i:i+BATCH]; xb,yb=x[ix],y[ix]
                c=carrier_tensor(len(ix),amp,'cpu'); xp=xb.view(-1,8,3).clone(); xp[:,:,2:3]+=c
                pp=model(xb); pc=model(xp.reshape(len(ix),-1)); td=nominal_delta(xb,amp,c[:,:,0])
                loss=nn.functional.mse_loss(pp,yb)+GAMMA*nn.functional.mse_loss(pc-pp,td[:,None])
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): metric=float(((model(ds['xte'])-ds['yte'])**2).mean())
        return metric, model, 'cpu'

def main():
    ds0 = get_dataset('dynamics', 0, n_train=400, n_test=200)
    # Canonical baseline sweep; all candidate learning rates are included here.
    def base_fn(cfg):
        def run(seed):
            seed_all(seed); ds=get_dataset('dynamics', seed, 400, 200)
            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
    base_block=sweep_baseline(base_fn, [{'lr':v} for v in LR_GRID])
    best_lr=base_block['best_cfg']['lr']
    idea_blocks=[]
    for amp in AMP_GRID:
        r=evaluate(lambda seed, a=amp: train_carrier(get_dataset('dynamics', seed, 400, 200), EPOCHS, best_lr, a, seed)[0], SEEDS)
        idea_blocks.append({'amp':amp,'result':r})
    best=min(idea_blocks, key=lambda q:q['result']['mean'])
    # Re-train best configuration to retain models for an NN-scale behavior signature.
    sig=[]
    for seed in SEEDS:
        ds=get_dataset('dynamics', seed, 400, 200)
        metric, model, dev=train_carrier(ds,EPOCHS,best_lr,best['amp'],seed)
        with torch.no_grad():
            x=ds['xte'][:128].to(dev); c=carrier_tensor(len(x),best['amp'],dev)
            xp=x.view(-1,8,3).clone(); xp[:,:,2:3]+=c
            dp=(model(xp.reshape(len(x),-1))-model(x)).squeeze(1).cpu().numpy()
            actual=nominal_delta(x,best['amp'],c[:,:,0]).cpu().numpy()
        sig.append((float(np.mean(np.abs(dp))),float(np.mean(np.abs(actual))),float(np.corrcoef(dp,actual)[0,1])))
    sig_arr=np.asarray(sig)
    mechanism={'carrier_amp':best['amp'],'predicted_abs_delta_mean':float(sig_arr[:,0].mean()),'observed_nominal_abs_delta_mean':float(sig_arr[:,1].mean()),'predicted_observed_correlation_mean':float(sig_arr[:,2].mean()),'quadratic_prediction':'not confirmed at NN scale: only one amplitude was behaviorally tested','confirmed':False}
    report=make_report('dynamics','rnn_small',base_block,best['result'],{'mechanism_signature':mechanism,'idea_sweep':idea_blocks,'track_justification':'Dynamics matches the proposed hidden-state observability/reachability and stability/control structure.'})
    report['idea_sweep']=idea_blocks
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()