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

TRACK = 'oscillatory_poisson_dirichlet'
WIDTH, DEPTH = 64, 2

# Golden-ratio continued-fraction denominators: 1,1,2,3,5,8,13,21.
def denominators(j=8):
    out=[]; qm1, q = 0, 1
    for _ in range(j):
        out.append(q); qm1, q = q, q+qm1
    return np.asarray(out, dtype=np.float32)
Q = denominators(8)

def encode_np(x, kind):
    x = np.asarray(x, dtype=np.float32).reshape(-1, 1)
    if kind == 'baseline':
        return x
    z = 2*np.pi*x*Q.reshape(1,-1)
    # prescribed square-root energy normalization
    return np.concatenate([x, np.cos(z)/np.sqrt(Q), np.sin(z)/np.sqrt(Q)], axis=1).astype(np.float32)

class Net(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        layers=[]; d=input_dim
        for _ in range(DEPTH):
            layers += [nn.Linear(d, WIDTH), nn.ReLU()]; d=WIDTH
        layers.append(nn.Linear(d,1)); self.net=nn.Sequential(*layers)
    def forward(self,x): return self.net(x)

def make_train(kind, seed, cfg, keep=False):
    torch.manual_seed(10000+int(seed)); np.random.seed(10000+int(seed)); random.seed(10000+int(seed))
    d0 = get_dataset(TRACK, seed, 400, 400)
    d = dict(d0)
    for k in ('xtr','ytr','xte','yte'):
        d[k] = encode_np(d0[k], kind) if k.startswith('x') else np.asarray(d0[k], dtype=np.float32).reshape(-1,1)
    d['input_shape'] = d['xtr'].shape[1:]
    for k in ('xtr','ytr','xte','yte'):
        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
    model = Net(d['xtr'].shape[1])
    net, metric, hist = train_model(model, d, epochs=int(cfg['epochs']), lr=float(cfg['lr']), batch=128, log=lambda *_: None)
    if net is None: raise RuntimeError('training failed')
    if keep: return float(metric), net, d
    return float(metric)

def factory(kind, cfg):
    return lambda seed: make_train(kind, seed, cfg)

def math_check():
    # Verify recursion, absolute summability, and finite encoded feature energy.
    qs=denominators(16)
    recursion=bool(np.all(qs[2:] == qs[1:-1] + qs[:-2]))
    inv=float(np.sum(1/qs))
    x=np.linspace(0,1,1001)
    z=2*np.pi*x[:,None]*qs[None,:]
    energy=np.mean((np.cos(z)/np.sqrt(qs))**2+(np.sin(z)/np.sqrt(qs))**2,axis=0)
    return {'q_first_8':qs[:8].astype(int).tolist(),'recursion_ok':recursion,
            'sum_inverse_q_16':inv,'pair_energy_mean':float(energy.mean()),
            'pair_energy_range':[float(energy.min()),float(energy.max())],
            'confirmed': recursion and inv < 4.0 and bool(np.allclose(energy,1/qs,rtol=.03,atol=.01))}

def signature(base_cfg, idea_cfg):
    rows=[]
    for s in range(8):
        bm,bn,bd=make_train('baseline',s,base_cfg,True)
        im,inn,idata=make_train('idea',s,idea_cfg,True)
        dev=next(bn.parameters()).device
        xb_base=torch.as_tensor(bd['xte'],dtype=torch.float32, device=dev)
        xb_idea=torch.as_tensor(idata['xte'],dtype=torch.float32, device=next(inn.parameters()).device)
        y=bd['yte'].numpy().ravel()
        with torch.no_grad():
            bp=bn(xb_base).cpu().numpy().ravel()
            ip=inn(xb_idea).cpu().numpy().ravel()
        # measured high-frequency projection on the q=21 channel
        phase=np.sin(2*np.pi*np.asarray(bd['xte']).reshape(-1)*21); phase-=phase.mean()
        def proj(v):
            v=v-v.mean(); return float(np.dot(v,phase)/np.dot(phase,phase))
        rows.append({'seed':s,'observed_projection':proj(y),'baseline_projection':proj(bp),'idea_projection':proj(ip),'baseline_mse':bm,'idea_mse':im})
    obs=np.array([r['observed_projection'] for r in rows]); ba=np.array([r['baseline_projection'] for r in rows]); ia=np.array([r['idea_projection'] for r in rows])
    be=float(np.mean(np.abs(ba-obs))); ie=float(np.mean(np.abs(ia-obs)))
    return {'prediction':'lacunary features should represent the observed q=21 oscillatory projection more accurately than raw coordinates',
            'observed_projection_mean':float(obs.mean()),'baseline_projection_mean':float(ba.mean()),'idea_projection_mean':float(ia.mean()),
            'baseline_abs_projection_error':be,'idea_abs_projection_error':ie,'confirmed':bool(ie<be),'per_seed':rows}

def main():
    check=math_check()
    # Union parity: all idea settings are included in baseline sweep.
    grid=[{'lr':1e-3,'epochs':40},{'lr':3e-3,'epochs':60},{'lr':1e-2,'epochs':40}]
    base=sweep_baseline(lambda cfg: factory('baseline',cfg), grid)
    best=base['best_cfg']
    idea_runs=[evaluate(factory('idea',cfg)) for cfg in grid]
    idx=int(np.argmin([r['mean'] for r in idea_runs])); idea=idea_runs[idx]; chosen=grid[idx]
    # paired comparison uses same chosen config on baseline and idea, both all 8 seeds
    baseline_full=evaluate(factory('baseline',chosen))
    sig=signature(chosen,chosen)
    rep=make_report(TRACK,'mlp_tiny',base,idea,{'core_math':check,**sig})
    rep['baseline_full_at_idea_cfg']=baseline_full
    rep['idea_config_candidates']=grid; rep['chosen_idea_cfg']=chosen
    rep['track_justification']='Registered oscillatory elliptic PDE field: the task contains the high-frequency coordinate structure targeted by the embedding.'
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()