import sys, json, time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

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

# Custom matched structure: each example is a sequence of token distributions,
# and the target depends on the robust location of the clean token population.
META = {'name':'distribution_token_sequence','domain':'sequence',
        'description':'Regression from a sequence of quantile-vector token distributions; some training/test tokens are grossly corrupted.'}

def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    nt, m = 12, 16
    def make(n, train=False):
        y = rng.uniform(-1, 1, n).astype(np.float32)
        # Token distributions have a shared location signal plus token noise.
        centers = y[:,None] + rng.normal(0, .08, (n,nt)).astype(np.float32)
        levels = np.linspace(-.45,.45,m,dtype=np.float32)
        q = centers[:,:,None] + levels[None,None,:] + rng.normal(0,.035,(n,nt,m)).astype(np.float32)
        # contamination is present only in the observed input; target remains y.
        rate = .30
        bad = rng.rand(n,nt) < rate
        q += bad[:,:,None] * rng.choice([-1.,1.], size=(n,nt,1)).astype(np.float32) * 3.0
        return q.astype(np.float32), y[:,None]
    xtr,ytr=make(n_train,True); xte,yte=make(n_test,False)
    return {'xtr':torch.tensor(xtr), 'ytr':torch.tensor(ytr),
            'xte':torch.tensor(xte), 'yte':torch.tensor(yte),
            'task':'regression','metric':'mse','input_shape':(nt,m),'out_dim':1,
            'track':'distribution_token_sequence'}

class PoolNet(nn.Module):
    def __init__(self, m=16, mode='mean', delta=0.8):
        super().__init__(); self.mode=mode; self.delta=delta
        self.token = nn.Sequential(nn.Linear(m,32), nn.ReLU(), nn.Linear(32,16))
        self.head = nn.Sequential(nn.Linear(16,32),nn.ReLU(),nn.Linear(32,1))
    def pool(self,q):
        # q [B,T,M]. Pointwise Huber barycenter, safeguarded by fixed bracket.
        if self.mode=='mean': return q.mean(1)
        z=q.mean(1); lo=q.min(1).values; hi=q.max(1).values
        for _ in range(8):
            r=z[:,None,:]-q
            psi=r.clamp(-self.delta,self.delta)
            active=(r.abs()<self.delta).float()
            score=psi.mean(1); cur=active.mean(1)
            newton=z-score/(cur+1e-5)
            ok=(newton>=lo)&(newton<=hi)&(cur>1e-5)
            z=torch.where(ok,newton,(lo+hi)/2)
            r=z[:,None,:]-q; score=r.clamp(-self.delta,self.delta).mean(1)
            lo=torch.where(score<0,z,lo); hi=torch.where(score>0,z,hi)
        return z.sort(dim=-1).values
    def forward(self,x): return self.head(self.pool(self.token(x)))

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

def run(cfg, mode, seed):
    seed_all(seed); d=get_dataset(seed,400,400)
    net=PoolNet(mode=mode,delta=cfg.get('delta',.8))
    _, metric, hist=train_model(net,d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *a,**k:None)
    return float(metric)

def main():
    # Union parity: every lr tried for Huber is also swept for arithmetic.
    grid=[{'lr':x,'epochs':12,'delta':d} for x in (0.001,0.003,0.006) for d in (0.5,0.8,1.2)]
    base_grid=[{'lr':x,'epochs':12,'delta':d} for x in (0.001,0.003,0.006) for d in (0.5,0.8,1.2)]
    base=sweep_baseline(lambda c: lambda s: run(c,'mean',s),base_grid)
    # Equal 3-config idea sweep at the baseline-selected lr plus nearby deltas.
    bestlr=base['best_cfg']['lr']
    idea_grid=[{'lr':bestlr,'epochs':12,'delta':d} for d in (0.5,0.8,1.2)]
    ir=[]
    for c in idea_grid: ir.append((c,evaluate(lambda s,c=c:run(c,'huber',s))))
    bestc, idea=min(ir,key=lambda z:z[1]['mean'])
    # Signature measured from independently trained benchmark systems.
    seed_all(123); ds_sig=get_dataset(123,400,400)
    bnet,_,_=train_model(PoolNet(mode='mean'),ds_sig,epochs=12,lr=base['best_cfg']['lr'],batch=128,log=lambda *a,**k:None)
    hnet,_,_=train_model(PoolNet(mode='huber',delta=bestc['delta']),ds_sig,epochs=12,lr=bestc['lr'],batch=128,log=lambda *a,**k:None)
    with torch.no_grad():
        x=torch.zeros(1,12,16); x[:,:,:]=torch.linspace(-.45,.45,16)
        def shift(net):
            clean=net.pool(x).mean(); xx=x.clone(); xx[:,0,:]+=10.0
            return float((net.pool(xx).mean()-clean).abs())
        changed=float(shift(hnet)); changedb=float(shift(bnet))
    sig={'prediction':'a gross token should shift Huber pooled quantiles less than arithmetic pooling',
         'predicted_bounded_score':bestc['delta']/12.,'observed_huber_shift':changed,
         'observed_mean_shift':changedb,'confirmed':bool(changed < changedb)}
    rep=make_report('distribution_token_sequence','custom_poolnet',base,idea,
        {'prediction':sig['prediction'],'predicted_bounded_score':sig['predicted_bounded_score'],
         'observed_huber_shift':changed,'observed_mean_shift':changedb,'confirmed':sig['confirmed'],
         'custom_track':{'name':'distribution_token_sequence','file':'hub_bench.py','domain':'sequence'},
         'idea_grid':[{**c,'mean':r['mean']} for c,r in ir], 'selected_idea_cfg':bestc})
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()