import sys, json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report

OUT = Path(__file__).with_name('bench_report.json')
EPOCHS = 18
BATCH = 128
RHO_AUG = 0.50
DENSITIES = [0.25, 0.5, 0.75]  # baseline real-space density; union with idea lr grid is shared
LRS = [1e-3, 3e-3, 1e-2]


def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'


def flatten_params(net):
    # Prune every affine weight tensor, excluding biases. Keep tensor shapes for forward.
    refs = []
    for mod in net.modules():
        if isinstance(mod, nn.Linear): refs.append(mod)
    return refs


def run_one(seed, lr, augmented, density=0.5):
    seed_all(seed)
    ds = get_dataset('tabular', seed=seed, n_train=400, n_test=200)
    # Make identical random base architecture/initialization for paired systems.
    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
    refs = flatten_params(net)
    dev = device()
    try:
        net = net.to(dev)
        xtr = torch.as_tensor(ds['xtr'], dtype=torch.float32, device=dev)
        ytr = torch.as_tensor(ds['ytr'], dtype=torch.float32, device=dev)
        xte = torch.as_tensor(ds['xte'], dtype=torch.float32, device=dev)
        yte = torch.as_tensor(ds['yte'], dtype=torch.float32, device=dev)
        # Freeze all weights; score parameters are the only trainable quantities.
        for p in net.parameters(): p.requires_grad_(False)
        real_scores = [nn.Parameter(torch.randn_like(m.weight, device=dev)) for m in refs]
        dummy_scores = [nn.Parameter(torch.randn_like(m.weight, device=dev)) for m in refs] if augmented else []
        params = real_scores + dummy_scores
        opt = torch.optim.Adam(params, lr=lr)
        # Exact fixed candidate quota. For a doubled layer, K=floor(rho_aug*2M).
        quotas = []
        for m in refs:
            M = m.weight.numel()
            quotas.append(max(1, min(2*M if augmented else M,
                                     int(math.floor((RHO_AUG if augmented else density) * (2*M if augmented else M))))))
        for epoch in range(EPOCHS):
            perm = torch.randperm(xtr.shape[0], device=dev)
            for start in range(0, xtr.shape[0], BATCH):
                ix = perm[start:start+BATCH]
                hlist = []
                for m, sr, sd, K in zip(refs, real_scores, dummy_scores if augmented else [None]*len(refs), quotas):
                    sa = torch.cat([sr.reshape(-1), sd.reshape(-1)]) if augmented else sr.reshape(-1)
                    hard = torch.zeros_like(sa); hard[torch.topk(sa, K).indices] = 1.
                    st = (hard - sa).detach() + sa
                    hlist.append(st[:m.weight.numel()].reshape_as(m.weight))
                # Functional forward with masked weights, preserving the shared MLP.
                z = xtr[ix]
                li = 0
                for mod in net.modules():
                    if isinstance(mod, nn.Linear):
                        z = torch.nn.functional.linear(z, mod.weight * hlist[li], mod.bias)
                        li += 1
                        if li < len(refs): z = torch.relu(z)
                loss = torch.mean((z - ytr[ix]) ** 2)
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            masks=[]; real_counts=[]; dummy_counts=[]
            for m, sr, sd, K in zip(refs, real_scores, dummy_scores if augmented else [None]*len(refs), quotas):
                sa = torch.cat([sr.reshape(-1), sd.reshape(-1)]) if augmented else sr.reshape(-1)
                inds = torch.topk(sa, K).indices
                mask = torch.zeros(m.weight.numel(), device=inds.device); real_ind = inds[inds < m.weight.numel()]; mask[real_ind] = 1.; masks.append(mask.reshape_as(m.weight))
                real_counts.append(int((inds < m.weight.numel()).sum().item()))
                dummy_counts.append(int(K - real_counts[-1]))
            z=xte; li=0
            for mod in net.modules():
                if isinstance(mod, nn.Linear):
                    z=torch.nn.functional.linear(z, mod.weight*masks[li], mod.bias); li+=1
                    if li<len(refs): z=torch.relu(z)
            metric=float(torch.mean((z-yte)**2).item())
        return metric, {'real_density': float(np.mean([r/m.weight.numel() for r,m in zip(real_counts,refs)])),
                        'dummy_fraction': float(sum(dummy_counts)/sum(quotas)), 'real_counts': real_counts}
    except Exception as e:
        if dev == 'cuda':
            torch.cuda.empty_cache()
            # deterministic CPU retry
            return run_one_cpu(seed, lr, augmented, density)
        raise


def run_one_cpu(seed, lr, augmented, density=0.5):
    old=torch.cuda.is_available
    # Explicit CPU implementation by temporarily using a compact equivalent call is avoided;
    # rerun after forcing tensors/model to CPU through the shared function's code path.
    # In normal bench environments CUDA or CPU path above succeeds; this fallback is a direct small model.
    seed_all(seed); ds=get_dataset('tabular', seed=seed, n_train=400, n_test=200)
    net=make_model('mlp_tiny', ds['input_shape'], ds['out_dim']); refs=flatten_params(net)
    xtr=torch.tensor(ds['xtr'],dtype=torch.float32); ytr=torch.tensor(ds['ytr'],dtype=torch.float32)
    xte=torch.tensor(ds['xte'],dtype=torch.float32); yte=torch.tensor(ds['yte'],dtype=torch.float32)
    for p in net.parameters(): p.requires_grad_(False)
    sr=[nn.Parameter(torch.randn_like(m.weight)) for m in refs]; sd=[nn.Parameter(torch.randn_like(m.weight)) for m in refs] if augmented else []
    opt=torch.optim.Adam(sr+sd,lr=lr); qs=[max(1,int(math.floor((RHO_AUG if augmented else density)*(2*m.weight.numel() if augmented else m.weight.numel())))) for m in refs]
    for _ in range(EPOCHS):
      for st in range(0,len(xtr),BATCH):
        hs=[]
        for m,a,b,k in zip(refs,sr,sd if augmented else [None]*len(refs),qs):
          q=torch.cat([a.flatten(),b.flatten()]) if augmented else a.flatten(); hard=torch.zeros_like(q); hard[torch.topk(q,k).indices]=1; hs.append(((hard-q).detach()+q)[:m.weight.numel()].reshape_as(m.weight))
        z=xtr[st:st+BATCH]; j=0
        for mod in net.modules():
          if isinstance(mod,nn.Linear): z=nn.functional.linear(z,mod.weight*hs[j],mod.bias); j+=1; z=torch.relu(z) if j<len(refs) else z
        loss=((z-ytr[st:st+BATCH])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
      hs=[]; rc=[]; dc=[]
      for m,a,b,k in zip(refs,sr,sd if augmented else [None]*len(refs),qs):
        q=torch.cat([a.flatten(),b.flatten()]) if augmented else a.flatten(); ind=torch.topk(q,k).indices; mask=torch.zeros(m.weight.numel()); real_ind=ind[ind<m.weight.numel()]; mask[real_ind]=1.; hs.append(mask.reshape_as(m.weight)); rc.append(int(real_ind.numel())); dc.append(k-rc[-1])
      z=xte;j=0
      for mod in net.modules():
       if isinstance(mod,nn.Linear): z=nn.functional.linear(z,mod.weight*hs[j],mod.bias);j+=1;z=nn.functional.relu(z) if j<len(refs) else z
      return float(((z-yte)**2).mean()), {'real_density':float(np.mean([r/m.weight.numel() for r,m in zip(rc,refs)])),'dummy_fraction':float(sum(dc)/sum(qs)),'real_counts':rc}


def main():
    # Baseline method knobs (density and lr) and all idea lrs share the same union.
    grid=[{'lr':lr,'density':rho} for lr in LRS for rho in DENSITIES]
    def mk(cfg): return lambda s: run_one(s,cfg['lr'],False,cfg['density'])[0]
    base=sweep_baseline(mk,grid)
    # Three idea settings: baseline-best lr plus two nearby values, all in baseline grid.
    idea_lrs=sorted(set(LRS))
    idea_trials=[]
    for lr in idea_lrs:
        r=evaluate(lambda s, lr=lr: run_one(s,lr,True)[0], seeds=(0,1,2,3))
        idea_trials.append({'cfg':{'lr':lr,'rho_aug':RHO_AUG},'mean':r['mean']})
    best_idea_lr=min(idea_trials,key=lambda x:x['mean'])['cfg']['lr']
    idea=evaluate(lambda s: run_one(s,best_idea_lr,True)[0])
    sig=[run_one(s,best_idea_lr,True)[1] for s in range(8)]
    sig_b=[run_one(s,base['best_cfg']['lr'],False,base['best_cfg']['density'])[1] for s in range(8)]
    obs_dummy=float(np.mean([x['dummy_fraction'] for x in sig]))
    obs_density=float(np.mean([x['real_density'] for x in sig]))
    signature={'prediction':'For iid initialized real/dummy scores and fixed Top-K, expected dummy fraction is 0.5 and expected real density is rho_aug=0.5.',
      'predicted_dummy_fraction':0.5,'observed_dummy_fraction_mean':obs_dummy,
      'predicted_real_density':RHO_AUG,'observed_real_density_mean':obs_density,
      'baseline_observed_real_density':float(np.mean([x['real_density'] for x in sig_b])),
      'confirmed':bool(abs(obs_dummy-.5)<.12 and abs(obs_density-RHO_AUG)<.12)}
    rep=make_report('tabular','mlp_tiny',base,idea,{'track_choice':'tabular matches pruning/regularization of a frozen MLP; identical MLP and task are used for both systems.','shared_lr_grid':LRS,'baseline_density_grid':DENSITIES,'idea_sweep':idea_trials,'idea_best_cfg':{'lr':best_idea_lr,'rho_aug':RHO_AUG},'signature':signature})
    rep['bench_report']=rep.copy()
    rep['idea']['sweep']=idea_trials
    OUT.write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))

if __name__=='__main__': main()