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

SEED = 2235
WIDTH = 64
HIDDEN = 16
NTR = 800
NTE = 400
SEEDS = tuple(range(8))
# Smooth activation used by the standard branch MLP; its first two Taylor
# terms produce the stated quadratic Gram module.
def phi(s):
    return s + 0.5*s*s + 0.1*s*s*s

def math_check():
    torch.manual_seed(SEED)
    n, d, h = 7, WIDTH, HIDDEN
    x = torch.randn(19, n, dtype=torch.float64)
    base = torch.randn(n+h, d, dtype=torch.float64)
    def full(W):
        u, v = W[:n], W[n:]
        return (phi(x @ u) @ v.T)
    def quad(W):
        u, v = W[:n], W[n:]
        # Tr(WW^T A_j(x)) = sum_i v_ij (x^T u_i)
        return ((x @ u) @ v.T)
    es = np.logspace(-3, -0.7, 8)
    rem = np.array([float((full(e*base)-quad(e*base)).abs().mean()) for e in es])
    slope = float(np.polyfit(np.log(es), np.log(rem), 1)[0])
    q, _ = torch.linalg.qr(torch.randn(d,d, dtype=torch.float64))
    gram_err = float((quad(base@q)-quad(base)).abs().max())
    return {'epsilon_values': es.tolist(), 'remainder_values': rem.tolist(),
            'observed_remainder_slope': slope, 'predicted_remainder_slope': 3.0,
            'gram_invariance_max_error': gram_err,
            'prediction_pass': bool(abs(slope-3.0)<0.15 and gram_err<1e-5)}

class BranchMLP(nn.Module):
    def __init__(self, n, idea=False):
        super().__init__(); self.idea=idea
        self.W=nn.Parameter(0.05*torch.randn(n+HIDDEN, WIDTH))
        self.proj=nn.Linear(HIDDEN, 1)
    def forward(self, x):
        u,v=self.W[:x.shape[-1]],self.W[x.shape[-1]:]
        a=x @ u
        z = (a @ v.T) if self.idea else (phi(a) @ v.T)
        return self.proj(z)

def make_train_fn(cfg, idea):
    def run(seed):
        torch.manual_seed(SEED + 1000 + int(seed)); np.random.seed(SEED+int(seed)); random.seed(SEED+int(seed))
        d=get_dataset('tabular', int(seed), n_train=NTR, n_test=NTE)
        m=BranchMLP(d['input_shape'][0], idea=idea)
        _, metric, _ = train_model(m, d, epochs=30, lr=float(cfg['lr']), batch=128, weight_decay=float(cfg['weight_decay']), log=lambda *_: None)
        return float(metric)
    return run

def trained_signature():
    # Obtain a trained baseline model and re-test the cubic remainder prediction
    # on its actual learned branch weights, rather than on a synthetic graph.
    seed=0; d=get_dataset('tabular', seed, n_train=NTR, n_test=NTE)
    torch.manual_seed(SEED+1000); m=BranchMLP(d['input_shape'][0], idea=False)
    m,_,_=train_model(m,d,epochs=30,lr=0.003,batch=128,weight_decay=0.0,log=lambda *_: None)
    m = m.cpu(); x=d['xte'][:128].cpu(); W=m.W.detach().cpu(); n=x.shape[1]; u,v=W[:n],W[n:]
    def f(ww):
        uu,vv=ww[:n],ww[n:]; return (phi(x@uu)@vv.T) @ m.proj.weight.detach().T + m.proj.bias.detach()
    def q(ww):
        uu,vv=ww[:n],ww[n:]; return ((x@uu)@vv.T) @ m.proj.weight.detach().T + m.proj.bias.detach()
    es=np.logspace(-2.5,-0.8,7); vals=np.array([float((f(e*W)-q(e*W)).abs().mean()) for e in es])
    slope=float(np.polyfit(np.log(es),np.log(np.maximum(vals,1e-30)),1)[0])
    perm=torch.randperm(W.shape[1]); perr=float((f(W[:,perm])-f(W)).abs().max())
    return {'source':'trained baseline model on tabular test inputs','epsilon_values':es.tolist(),
            'observed_remainder_slope':slope,'predicted_remainder_slope':3.0,
            'permutation_max_error':perr,
            'confirmed':bool(abs(slope-3.0)<0.25 and perr<1e-5)}

def main():
    checks=math_check(); print('MATH',json.dumps(checks))
    grid=[{'lr':0.001,'weight_decay':0.0},{'lr':0.003,'weight_decay':0.0},{'lr':0.01,'weight_decay':0.0}]
    base=sweep_baseline(lambda c: make_train_fn(c,False), grid, seeds=(0,1,2,3))
    idea_trials=[]
    for cfg in grid:
        r=evaluate(make_train_fn(cfg,True), SEEDS)
        idea_trials.append({'cfg':cfg,'result':r})
    best=min(idea_trials,key=lambda z:z['result']['mean'])
    idea={'best_cfg':best['cfg'],'sweep':[{'cfg':z['cfg'],'mean':z['result']['mean']} for z in idea_trials],
          'per_seed':best['result']['per_seed'],'mean':best['result']['mean'],'std':best['result']['std'],'n':best['result']['n']}
    sig=trained_signature()
    rep=make_report('tabular','custom_branch_mlp',base,idea,{'math_check':checks,'trained_model':sig,
        'custom_track':None})
    rep['protocol_notes']={'paired_seeds':list(SEEDS),'n_train':NTR,'n_test':NTE,
        'architecture':'same width-64 exchangeable branch MLP; only phi vs quadratic Taylor block differs',
        'baseline_grid':grid,'idea_grid':grid}
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()