import sys, json, random
from pathlib import Path
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import train_model, evaluate, sweep_baseline, make_report
from bench.custom_tracks import burgers_periodic_split as track

META = {'name':'burgers_front_tokens_bench','domain':'pde','description':'Burgers one-step regression with analytical RH front tokens versus matched CNN.'}

# Shared base architecture: both systems have the same hidden convolutional encoder.
class Encoder(nn.Module):
    def __init__(self, n=32, width=32):
        super().__init__(); self.net=nn.Sequential(nn.Conv1d(1,width,5,padding=2),nn.ReLU(),nn.Conv1d(width,width,5,padding=2),nn.ReLU(),nn.Conv1d(width,width,3,padding=1),nn.ReLU())
    def forward(self,x): return self.net(x)

class Baseline(nn.Module):
    def __init__(self,n=32):
        super().__init__(); self.enc=Encoder(n); self.head=nn.Conv1d(32,1,1)
    def forward(self,x): return self.head(self.enc(x))

def front_info(x, threshold=.18):
    # Differentiable-free detector is intentionally input preprocessing; traces are noisy local estimates.
    u=x[:,0,:]; d=(u[:,:,None] if False else None)
    jumps=(u[:,1:]-u[:,:-1]).abs(); B,N=u.shape
    mask=jumps > threshold
    # one token: strongest detected edge (periodic edge included)
    vals, idx=jumps.max(dim=1); idx=idx.clamp(1,N-2)
    grid=torch.arange(N,device=x.device).float()[None,:]
    pos=(idx.float()+.5)/(N-1)
    il=idx-1; ir=idx+1
    ul=u.gather(1,il[:,None]).squeeze(1); ur=u.gather(1,ir[:,None]).squeeze(1)
    active=(vals>threshold).float()
    return pos,ul,ur,active

class FrontTokens(nn.Module):
    def __init__(self,n=32):
        super().__init__(); self.enc=Encoder(n)
        # residual head uses same parameter count scale, while traces are estimated from local observed grid values
        self.res=nn.Conv1d(32,1,1)
        self.trace=nn.Sequential(nn.Linear(32,16),nn.Tanh(),nn.Linear(16,2))
        self.n=n
    def forward(self,x):
        h=self.enc(x); residual=self.res(h)
        pos,ul,ur,active=front_info(x)
        # RH speed for Burgers: (F(ur)-F(ul))/(ur-ul) = (ul+ur)/2.
        speed=.5*(ul+ur)
        newpos=(pos + .08*speed).remainder(1.0)
        # Render a periodic step from estimated traces; trace MLP receives local encoded context.
        b=torch.arange(x.shape[0],device=x.device)
        ctx=h[b,:,torch.clamp((pos*(self.n-1)).long(),0,self.n-1)]
        delta=self.trace(ctx); aul=ul+0.15*delta[:,0]; aur=ur+0.15*delta[:,1]
        xx=torch.arange(self.n,device=x.device).float()[None,:]/(self.n-1)
        # smooth sigmoid rendering keeps gradients through positions and traces
        dist=((xx-newpos[:,None]+.5)%1.0)-.5
        front=aur[:,None]+(aul-aur)[:,None]*torch.sigmoid(-dist/0.035)
        return residual + active[:,None,None]*front[:,None,:], (pos,newpos,ul,ur,speed,active)

def idea_forward(net,x): return net(x)[0]
def ds(seed):
    d=track.get_dataset(seed, n_train=400, n_test=160)
    for k in ('xtr','ytr','xte','yte'):
        d[k]=torch.as_tensor(d[k], dtype=torch.float32)
    return d
def seed_all(s): random.seed(s); np.random.seed(s); torch.manual_seed(s)
def train_one(kind,cfg,seed,return_net=False):
    seed_all(seed); d=ds(seed); net=Baseline() if kind=='base' else FrontTokens()
    if kind=='base': out=train_model(net,d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *_:None)
    else:
        # Official trainer expects tensor output; wrap the idea system without changing its learned system.
        class Wrap(nn.Module):
            def __init__(self,m): super().__init__(); self.m=m
            def forward(self,z): return self.m(z)[0]
        out=train_model(Wrap(net),d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *_:None)
    if return_net: return out,out[0],d
    return out[1] if out[1] is not None else float('nan')

def signature():
    pred=[]; obs=[]
    try:
        for s in range(8):
            out,net,d=train_one('idea',{'lr':.003,'epochs':15},s,True)
            if out[0] is None: continue
            net.cpu(); net.eval()
            x=torch.as_tensor(d['xte'], dtype=torch.float32); y=torch.as_tensor(d['yte'], dtype=torch.float32)
            with torch.no_grad(): z,info=net.m(x); pos,np_,ul,ur,sp,act=info
            def cross(a):
                a=a[:,0,:] if a.ndim==3 else a
                q=(a[:,:-1]>=0.5)&(a[:,1:]<0.5); ix=q.float().argmax(1); return (ix.float()+.5)/31
            target=cross(y); observed=target-cross(x); m=act>0
            if m.any(): pred += ((np_-pos)[m]).tolist(); obs += observed[m].tolist()
    except Exception as e:
        return {'predicted_mean_speed_displacement':float('nan'),'observed_mean_displacement':float('nan'),'relative_error':float('nan'),'confirmed':False,'error':str(e)[:180]}
    if not pred: return {'predicted_mean_speed_displacement':float('nan'),'observed_mean_displacement':float('nan'),'relative_error':float('nan'),'confirmed':False}
    a,b=float(np.mean(pred)),float(np.mean(obs)); rel=abs(a-b)/(abs(b)+1e-6)
    return {'predicted_mean_speed_displacement':a,'observed_mean_displacement':b,'relative_error':rel,'confirmed':bool(rel<.15)}

def main():
    # Union-parity grid: all idea learning rates are also baseline-evaluated.
    grid=[{'lr':lr,'epochs':ep} for lr in (.0015,.003,.006) for ep in (15,)]
    base=sweep_baseline(lambda c: lambda s:train_one('base',c,s),grid)
    idea_grid=grid
    idea_tuning=[]
    for cfg in idea_grid:
        r=evaluate(lambda s,c=cfg:train_one('idea',c,s), seeds=(0,1,2,3))
        idea_tuning.append({'cfg':cfg,'mean':r['mean']})
    idea_cfg=min(idea_tuning,key=lambda z:z['mean'])['cfg']
    idea=evaluate(lambda s:train_one('idea',idea_cfg,s))
    idea['selected_cfg']=idea_cfg
    extra={'idea_sweep':idea_tuning,'custom_track':{'name':'burgers_front_tokens_bench','file':'bench_front_tokens.py','domain':'pde'},'mechanism_signature':signature()}
    rep=make_report('burgers_periodic_split','cnn_shared',base,idea,extra)
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()