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

SEEDS = tuple(range(8)); SWEEP_SEEDS = tuple(range(4))
NTR, NTE, EPOCHS, FT_EPOCHS = 400, 200, 22, 5
LRS = [1e-3, 3e-3, 1e-2]
KEEP = 0.50

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

def dataset(seed):
    d = get_dataset('tabular', seed=seed, n_train=NTR, n_test=NTE)
    # train_model expects tensors; preserve all metadata
    return {**d, 'xtr': torch.as_tensor(d['xtr'], dtype=torch.float32),
            'ytr': torch.as_tensor(d['ytr'], dtype=torch.float32),
            'xte': torch.as_tensor(d['xte'], dtype=torch.float32),
            'yte': torch.as_tensor(d['yte'], dtype=torch.float32)}

def linear_layers(net):
    return [m for m in net.modules() if isinstance(m, nn.Linear)]

def masks_for(net, x, method):
    ls = linear_layers(net); h = x.to(ls[0].weight.device); masks=[]
    for i, layer in enumerate(ls[:-1]):
        with torch.no_grad(): z = layer(h); h = torch.relu(z)
        n = z.shape[1]; k=max(1, int(round(n*KEEP)))
        outnorm = ls[i+1].weight.detach().cpu().numpy().T
        if method == 'magnitude': score=np.linalg.norm(outnorm, axis=1)
        else:
            a=z.detach().cpu().numpy()
            lo=np.quantile(a, .01, axis=0); hi=np.quantile(a, .99, axis=0)
            visible=(lo <= 0.0) & (hi >= 0.0)
            used=np.linalg.norm(outnorm, axis=1) > 1e-8
            # visibility is the structural priority; downstream norm breaks ties
            score=visible.astype(np.float64)*1e6 + np.linalg.norm(outnorm,axis=1)
            score=np.where(used, score, -1e6)
        idx=np.argsort(score)[-k:]
        m=np.zeros(n, dtype=np.float32); m[idx]=1; masks.append(m)
    return masks

def apply_masks(net, masks):
    ls=linear_layers(net)
    with torch.no_grad():
        for i,m in enumerate(masks):
            t=torch.as_tensor(m, device=ls[i].weight.device)
            ls[i].weight.mul_(t[:,None]); ls[i].bias.mul_(t)
            ls[i+1].weight.mul_(t[None,:])

def run(method, lr, seed, want_signature=False):
    seed_all(seed); d=dataset(seed)
    # This is the same base architecture and same optimizer path on both sides.
    net=make_model('mlp_tiny', d['input_shape'], d['out_dim'])
    net, _, _=train_model(net,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *a,**k:None)
    masks=masks_for(net,d['xtr'],method); apply_masks(net,masks)
    net, metric, _=train_model(net,d,epochs=FT_EPOCHS,lr=lr,batch=128,log=lambda *a,**k:None)
    if want_signature: return float(metric), net, d, masks
    return float(metric)

def factory(method, cfg):
    return lambda seed: run(method, float(cfg['lr']), seed)

def mechanism_signature():
    metric, net, d, masks=run('task_visible', 3e-3, 0, True)
    ls=linear_layers(net); x=d['xte'].to(ls[0].weight.device)
    predicted=[]; observed=[]
    # On each trained model, test the local affine/zero claim for selected axes.
    for i, m in enumerate(masks):
        with torch.no_grad():
            h=x
            for q in range(i): h=torch.relu(ls[q](h))
            z=ls[i](h); lo=torch.quantile(z,.01,dim=0); hi=torch.quantile(z,.99,dim=0)
        one=((lo>0)|(hi<0)).cpu().numpy() & (m>0)
        if not one.any(): continue
        # Compare each retained one-sided unit's ReLU contribution with its affine replacement.
        w=ls[i].weight.detach(); b=ls[i].bias.detach(); out=ls[i+1].weight.detach()
        zz=z.detach()
        for j in np.flatnonzero(one):
            pred=torch.relu(zz[:,j]); repl=zz[:,j] if float(lo[j])>0 else torch.zeros_like(zz[:,j])
            predicted.append(float(torch.max(torch.abs(pred-repl))))
            observed.append(float(torch.linalg.vector_norm(out[:,j]).cpu()))
    maxerr=max(predicted) if predicted else 0.0
    return {'prediction':'strictly one-sided selected ReLU axes are affine or zero on the test patch',
            'predicted_max_replacement_error':0.0, 'observed_max_replacement_error':maxerr,
            'tested_one_sided_axes':len(predicted), 'downstream_norm_sum':float(sum(observed)),
            'confirmed':bool(maxerr < 1e-5)}

def main():
    grid=[{'lr':x} for x in LRS]
    base=sweep_baseline(lambda c: factory('magnitude',c), grid, seeds=SWEEP_SEEDS)
    idea_trials=[{'cfg':c,'result':evaluate(factory('task_visible',c),SEEDS)} for c in grid]
    best=min(idea_trials,key=lambda q:q['result']['mean'])
    rep=make_report('tabular','mlp_tiny',base,best['result'],{
      'track_justification':'ReLU hidden-axis pruning is structurally matched to the tabular MLP track; both systems use the identical trained mlp_tiny and optimizer budgets.',
      'pruning_fraction':KEEP, 'idea_config':best['cfg'], 'idea_sweep':idea_trials,
      'parameter_fraction_estimate':KEEP*KEEP, 'mechanism_signature':mechanism_signature()})
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()