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

# Sequence is structurally matched: the target depends on correlations throughout a window.
# We replace the independently parameterized transformer encoder with one shared latent rule.
class AdaptiveSharedSequence(nn.Module):
    def __init__(self, win, d=64, tmax=6, ponder=0.001):
        super().__init__()
        self.win, self.d, self.tmax, self.ponder = win, d, tmax, ponder
        self.inp = nn.Linear(1, d)
        self.pos = nn.Parameter(torch.zeros(1, win, d))
        nn.init.normal_(self.pos, std=.02)
        self.norm = nn.LayerNorm(d)
        self.rule = nn.Sequential(nn.Linear(d, 128), nn.GELU(), nn.Linear(128, d))
        self.halt = nn.Linear(d, 1)
        self.head = nn.Linear(win*d, 1)

    def forward(self, x, return_aux=False):
        s = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
        acc = torch.zeros_like(s)
        mass = torch.zeros(x.shape[0], 1, device=x.device)
        steps = torch.zeros_like(mass)
        for _ in range(self.tmax):
            s = s + 0.20 * self.rule(self.norm(s))
            h = torch.sigmoid(self.halt(s.mean(dim=1)))
            delta = torch.minimum(h, 1.0-mass)
            acc = acc + delta.unsqueeze(-1) * s
            mass = mass + delta
            steps = steps + (mass < 1.0-1e-3).float()
        acc = acc + (1.0-mass).unsqueeze(-1)*s
        out = self.head(acc.reshape(x.shape[0], -1))
        if return_aux:
            return out, steps.squeeze(1), mass.squeeze(1)
        return out

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 run_one(kind, seed, lr, epochs):
    seed_all(seed)
    ds = get_dataset('sequence', seed, n_train=400, n_test=200)
    if kind == 'base':
        net = make_model('transformer_tiny', ds['input_shape'], ds['out_dim'])
    else:
        net = AdaptiveSharedSequence(ds['input_shape'][0], tmax=6)
    net, metric, hist = train_model(net, ds, epochs=epochs, lr=lr, batch=128)
    if net is None: return float('nan'), {}
    net.eval()
    with torch.no_grad():
        dev = next(net.parameters()).device
        pred = net(ds['xte'].to(dev))
        if isinstance(pred, tuple): pred = pred[0]
        mse = torch.mean((pred.cpu()-ds['yte'].cpu())**2).item()
        aux={}
        if kind == 'idea':
            _, st, mass = net(ds['xte'].to(dev), return_aux=True)
            aux={'avg_microsteps': float(st.mean()), 'mean_halt_mass': float(mass.mean()),
                 'pred_std': float(pred.std()), 'target_std': float(ds['yte'].std())}
        else:
            aux={'pred_std': float(pred.std()), 'target_std': float(ds['yte'].std())}
    return mse, aux

def fn(kind, lr, epochs):
    return lambda seed: run_one(kind, seed, lr, epochs)[0]

def main():
    # Shared union: baseline and idea both evaluated at every lr in the 3-point grid.
    grid=[{'lr':1e-3,'epochs':18},{'lr':3e-3,'epochs':18},{'lr':6e-3,'epochs':18}]
    base=sweep_baseline(lambda c: fn('base',c['lr'],c['epochs']), grid)
    # idea is run at all three settings, selecting by the same 4-seed sweep protocol
    idea_trials=[]
    for c in grid:
        r=evaluate(fn('idea',c['lr'],c['epochs']), seeds=(0,1,2,3))
        idea_trials.append({'cfg':c,'mean':r['mean']})
    best=min(idea_trials,key=lambda z:z['mean'])['cfg']
    idea_full=evaluate(fn('idea',best['lr'],best['epochs']))
    rep=make_report('sequence','transformer_tiny',base,idea_full,extra={
      'sweep_parity': {'union_grid':grid,'idea_sweep':idea_trials},
      'mechanism_signature': mechanism_signature(best)
    })
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))

def mechanism_signature(cfg):
    rows=[]
    for s in range(8):
        seed_all(s); ds=get_dataset('sequence',s,n_train=400,n_test=200)
        net=AdaptiveSharedSequence(ds['input_shape'][0],tmax=6)
        net,_,_=train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128)
        net.eval()
        with torch.no_grad():
            dev = next(net.parameters()).device
            p,st,m=net(ds['xte'].to(dev),return_aux=True)
        rows.append({'seed':s,'observed_steps':float(st.mean()),'halt_mass':float(m.mean()),'pred_std':float(p.std()),'target_std':float(ds['yte'].std())})
    observed=float(np.mean([r['observed_steps'] for r in rows]))
    # Engineering prediction: adaptive recurrence should terminate below Tmax on average.
    predicted=6.0
    return {'claim':'trained adaptive recurrence uses fewer than Tmax microsteps on average',
            'predicted_avg_steps_upper_bound':predicted,'observed_avg_steps':observed,
            'relative_reduction':float((predicted-observed)/predicted),
            'trained_model_measurements':rows,'confirmed':bool(observed < predicted-1e-6)}

if __name__=='__main__': main()