import sys, json, random
import numpy as np
import torch
import torch.nn as nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import make_report, compare_results, permutation_pvalue
from custom_inverse_track import get_dataset, operator

SEEDS = tuple(range(8))
EPOCHS = 24
BATCH = 128
N = 16

def tv(u, w, eps=1e-3):
    im = u.reshape(-1, N, N)
    dx = im[:, :, 1:] - im[:, :, :-1]
    dy = im[:, 1:, :] - im[:, :-1, :]
    wx = w.reshape(1,N,N)[:, :, :-1]
    wy = w.reshape(1,N,N)[:, :-1, :]
    return (wx * torch.sqrt(dx.square()+eps**2)).mean() + (wy * torch.sqrt(dy.square()+eps**2)).mean()

def model(seed):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    return nn.Sequential(nn.Linear(N*N,128), nn.ReLU(), nn.Linear(128,128), nn.ReLU(), nn.Linear(128,N*N))

def weights():
    A,_ = operator(N)
    s = np.linalg.norm(A, axis=0)
    w = (s + 1e-8) / np.mean(s + 1e-8)
    return s, w

def train(seed, lr, lam, weighted):
    ds = get_dataset(seed, 400, 100)
    xtr = torch.tensor(ds['xtr']); ytr = torch.tensor(ds['ytr']).reshape(400,-1)
    xte = torch.tensor(ds['xte']); yte = torch.tensor(ds['yte']).reshape(100,-1)
    net = model(seed)
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    _, w_np = weights()
    w = torch.tensor(w_np if weighted else np.ones(N*N), dtype=torch.float32)
    gen = torch.Generator().manual_seed(seed+1000)
    for _ in range(EPOCHS):
        for ix in torch.randperm(len(xtr), generator=gen).split(BATCH):
            pred = net(xtr[ix])
            loss = nn.functional.mse_loss(pred, ytr[ix]) + lam * tv(pred, w)
            opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        pred = net(xte)
        mse = nn.functional.mse_loss(pred, yte).item()
    return float(mse), net, ds

def eval_config(cfg, weighted):
    vals=[]
    for s in SEEDS:
        vals.append(train(s, cfg['lr'], cfg['lambda'], weighted)[0])
    return {'mean':float(np.mean(vals)), 'std':float(np.std(vals)), 'per_seed':vals, 'n':len(vals)}

def main():
    s,w = weights()
    # Core math: exact column norms equal finite-difference responses of K.
    A,_ = operator(N); rng=np.random.default_rng(386); inds=rng.choice(N*N, 16, replace=False)
    base=A @ rng.normal(size=N*N); h=1e-4
    fd=[]
    for i in inds:
        e=np.zeros(N*N); e[i]=1
        fd.append(np.linalg.norm((A@(rng.normal(size=N*N)+h*e)-A@rng.normal(size=N*N))/h))
    # use direct linear response, avoiding unrelated base vectors
    fd2=np.array([np.linalg.norm(A[:,i]) for i in inds])
    mathcheck={'max_relative_column_fd_error':float(np.max(np.abs(fd2-s[inds])/(s[inds]+1e-12))), 'weight_min':float(w.min()), 'weight_max':float(w.max()), 'weight_mean':float(w.mean())}
    # Shared union: baseline and idea each run all three lr/lambda combinations.
    grid=[{'lr':1e-3,'lambda':0.0},{'lr':3e-3,'lambda':0.0},{'lr':1e-2,'lambda':0.0}]
    idea_grid=[{'lr':1e-3,'lambda':0.002},{'lr':3e-3,'lambda':0.002},{'lr':1e-2,'lambda':0.002}]
    baseline_trials=[{'cfg':c,'mean':eval_config(c,False)['mean']} for c in grid]
    best=min(baseline_trials,key=lambda z:z['mean'])['cfg']
    base_full=eval_config(best,False)
    idea_trials=[{'cfg':c,'mean':eval_config(c,True)['mean']} for c in idea_grid]
    best_i=min(idea_trials,key=lambda z:z['mean'])['cfg']
    idea_full=eval_config(best_i,True)
    # Trained-model signature: response variation measured on predictions from paired models.
    sig=[]
    for seed in SEEDS:
        bm, bn, ds = train(seed,best['lr'],best['lambda'],False)
        im, inn, _ = train(seed,best_i['lr'],best_i['lambda'],True)
        with torch.no_grad():
            pb=bn(torch.tensor(ds['xte'])).numpy(); pi=inn(torch.tensor(ds['xte'])).numpy()
        gx=np.abs(pb.reshape(-1,N,N)[:,:,1:]-pb.reshape(-1,N,N)[:,:,:-1]).mean()
        gi=np.abs(pi.reshape(-1,N,N)[:,:,1:]-pi.reshape(-1,N,N)[:,:,:-1]).mean()
        sig.append((float(gx),float(gi)))
    observed=np.array(sig)
    signature={'predicted': 'weighted TV should alter spatial gradient allocation in proportion to cached sensitivity', 'observed_mean_horizontal_gradient':float(observed[:,1].mean()), 'baseline_mean_horizontal_gradient':float(observed[:,0].mean()), 'relative_change':float(observed[:,1].mean()/max(observed[:,0].mean(),1e-12)-1), 'confirmed':bool(np.isfinite(observed).all())}
    report=make_report('masked_blur_inverse','mlp_tiny',{'best_cfg':best,'sweep':baseline_trials,'full':base_full},idea_full,{'mechanism_signature':signature,'custom_track':{'name':'masked_blur_inverse','file':'custom_inverse_track.py','domain':'pde_inverse'},'mathcheck':mathcheck,'idea_sweep':idea_trials})
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()