import sys, json, random
from pathlib import Path
import numpy as np
import torch
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, sweep_baseline, make_report
from bench.protocol import evaluate
from bench.protocol import DEFAULT_SEEDS

SEEDS = tuple(range(8))
# Union of all lrs is used on both sides (search-space parity).
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 8
NTR, NTE = 400, 160

class AdaptiveChunkRNN(torch.nn.Module):
    """Same GRU/head as rnn_small, but schedules sequence tokens in chunks.

    A chunk boundary is selected from input-derived saturation (mean absolute
    angular velocity) and whale/load proxy (fraction of large controls). Hidden
    state is carried across chunks, so the intervention is scheduling only.
    """
    def __init__(self, base, cmin=2, cmax=8, ks=3.0, kw=2.0):
        super().__init__()
        self.rnn, self.head = base.rnn, base.head
        self.cmin, self.cmax, self.ks, self.kw = cmin, cmax, ks, kw
        self.last_trace = None

    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        # saturation normalized by a robust pendulum scale; whale proxy is
        # high-control occupancy, matching the serving controller's inputs.
        sat = (seq[:, :, 1].abs().mean(1) / 2.0).clamp(0, 1)
        whale = (seq[:, :, 2].abs() > 1.0).float().mean(1)
        chunks = (self.cmax - self.ks * sat - self.kw * whale).clamp(self.cmin, self.cmax).round().long()
        # Per-example chunk sizes are batched by using the minimum size in a
        # quantum; this is conservative and deterministic.
        c = int(chunks.min().item())
        h = None; states = []
        for j in range(0, seq.shape[1], c):
            out, h = self.rnn(seq[:, j:j+c], h)
            states.append(out[:, -1].detach())
        self.last_trace = {'chunks': chunks.detach(), 'states': states}
        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(): torch.cuda.manual_seed_all(seed)

def baseline_one(lr, seed, return_net=False):
    seed_all(seed); d=get_dataset('dynamics', seed, NTR, NTE)
    net=make_model('rnn_small', d['input_shape'], d['out_dim'])
    net, m, hist=train_model(net,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
    return (m, net, d) if return_net else m

def idea_one(lr, seed, return_net=False):
    seed_all(seed); d=get_dataset('dynamics', seed, NTR, NTE)
    base=make_model('rnn_small', d['input_shape'], d['out_dim'])
    net=AdaptiveChunkRNN(base)
    net, m, hist=train_model(net,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
    return (m, net, d) if return_net else m

def main():
    # Baseline sweep and idea sweep use identical lr union and 4 tuning seeds.
    grid=[{'lr':lr} for lr in LRS]
    base=sweep_baseline(lambda cfg: (lambda seed: baseline_one(cfg['lr'],seed)), grid=grid)
    # sweep_baseline expects a factory; inspect result shape defensively.
    best_cfg=base.get('best_cfg', grid[1])
    if isinstance(best_cfg, dict) and 'lr' in best_cfg: best_lr=best_cfg['lr']
    else: best_lr=3e-3
    idea_grid=[{'lr':lr} for lr in LRS]
    idea_sweep=sweep_baseline(lambda cfg: (lambda seed: idea_one(cfg['lr'],seed)), grid=idea_grid)
    # Explicit paired eight-seed results at the idea sweep winner.
    idea_lr=idea_sweep.get('best_cfg', {'lr':best_lr}).get('lr',best_lr)
    base_res=[baseline_one(idea_lr,s) for s in SEEDS]
    idea_res=[idea_one(idea_lr,s) for s in SEEDS]
    # Model-derived NN-scale mechanism signature: activation state ramps on
    # the trained systems, not an analytical identity.
    sig=[]
    for s in SEEDS:
        bm,bn,d=baseline_one(idea_lr,s,True); im,inn,_=idea_one(idea_lr,s,True)
        with torch.no_grad():
            # The training ladder may leave weights on a shared GPU.  The
            # signature is only a diagnostic, so run it robustly on CPU.
            bn=bn.cpu(); inn=inn.cpu(); xb=d['xte'][:128].cpu()
            old_cudnn=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
            bn.eval(); inn.eval()
            try:
                z=xb.view(len(xb),-1,3); out,h=bn.rnn(z); bstates=out
                _=inn(xb); astates=torch.stack(inn.last_trace['states'],1)
            finally:
                torch.backends.cudnn.enabled=old_cudnn
            br=float(torch.diff(bstates,dim=1).abs().mean())
            ar=float(torch.diff(astates,dim=1).abs().mean()) if astates.shape[1]>1 else 0.
        sig.append({'seed':s,'baseline_state_ramp':br,'adaptive_state_ramp':ar})
    br=float(np.mean([x['baseline_state_ramp'] for x in sig])); ar=float(np.mean([x['adaptive_state_ramp'] for x in sig]))
    report=make_report('dynamics','rnn_small',base,idea_sweep['full'],extra={
        'prediction':'adaptive chunking lowers high-load sequential-state ramp while preserving task function',
        'observed_baseline_state_ramp':br,'observed_adaptive_state_ramp':ar,
        'relative_reduction':1-ar/br if br else 0.,'confirmed': bool(ar < br),
        'note':'activation ramp is a hardware-independent proxy; no NVML power available'})
    report['custom_track']=None
    report['idea_sweep']=idea_sweep
    report['idea_lr']=idea_lr
    report['paired_baseline_at_idea_lr']=base_res
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()