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

SEEDS = tuple(range(8)); EPOCHS = 8
GRID = [{'lr': x, 'weight_decay': 0.0} for x in (0.0015, 0.003, 0.006)]

class FeatureNormCNN(nn.Module):
    # Identical cnn_small, differing only in the channel operation before its head.
    def __init__(self, out_dim, mode='bn', rho=.8, eps=1e-4):
        super().__init__(); base = cnn_small(out_dim)
        self.features = base.net[:9]       # through third ReLU, before Flatten
        self.tail = base.net[9:]           # exact shared Flatten/MLP/classifier
        self.mode, self.rho, self.eps = mode, rho, eps
        if mode == 'bn': self.norm = nn.BatchNorm1d(96)
        else:
            self.register_buffer('ema_R', torch.eye(96)); self.register_buffer('ema_mu', torch.zeros(96))
            self.register_buffer('initialized', torch.tensor(False))

    def _zca(self, z):
        if self.training:
            mu = z.mean(0); q = z - mu
            R = q.T @ q / max(1, z.shape[0]) + self.eps*torch.eye(96, device=z.device)
            with torch.no_grad():
                if not bool(self.initialized): self.ema_R.copy_(R.detach()); self.ema_mu.copy_(mu.detach()); self.initialized.fill_(True)
                else: self.ema_R.mul_(.95).add_(R.detach(), alpha=.05); self.ema_mu.mul_(.95).add_(mu.detach(), alpha=.05)
        else: mu, R = self.ema_mu, self.ema_R + self.eps*torch.eye(96, device=z.device)
        w,U = torch.linalg.eigh(R); A = (U*torch.rsqrt(torch.clamp(w,min=self.eps))) @ U.T
        return (z-mu) @ A

    def forward(self, x):
        h = self.features(x); b,c,hh,ww = h.shape
        z = h.permute(0,2,3,1).reshape(-1,c)
        z = self.norm(z) if self.mode == 'bn' else self._zca(z)
        h = z.reshape(b,hh,ww,c).permute(0,3,1,2)
        return self.tail(h)

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

def train_fn(mode, cfg):
    def run(seed):
        seed_all(seed); d=get_dataset('vision', seed, n_train=400, n_test=200)
        net=FeatureNormCNN(d['out_dim'], mode=mode, rho=cfg.get('rho',.8))
        _,metric,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],weight_decay=cfg['weight_decay'],batch=128,log=lambda *_:None)
        return metric
    return run

def signature(cfg):
    seed_all(0); d=get_dataset('vision',0,n_train=400,n_test=200)
    net=FeatureNormCNN(d['out_dim'],mode='zca',rho=cfg['rho']); net,_,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],weight_decay=cfg['weight_decay'],batch=128,log=lambda *_:None)
    net.eval()
    with torch.no_grad():
        h=net.features(d['xte']); z=h.permute(0,2,3,1).reshape(-1,96); y=net._zca(z)
        C=y.T@y/y.shape[0]; off=C-torch.diag(torch.diag(C))
        zs=(z-z.mean(0))/(z.std(0)+1e-6); ys=(y-y.mean(0))/(y.std(0)+1e-6)
        fid=torch.diag(zs.T@ys/zs.shape[0])
    observed_off=float(torch.linalg.norm(off)); observed_f=float(fid.min())
    return {'prediction':'decorrelation near zero with designated-channel fidelity >= rho_min', 'predicted_offdiag':0.0, 'predicted_min_fidelity_threshold':cfg['rho'], 'observed_offdiag':observed_off, 'observed_min_fidelity':observed_f, 'confirmed':bool(observed_off < .15 and observed_f >= cfg['rho']-.05)}

def main():
    base=sweep_baseline(lambda c:train_fn('bn',c),GRID,seeds=(0,1,2,3))
    idea_grid=[dict(c,rho=.8) for c in GRID]
    idea_runs=[]
    for c in idea_grid: idea_runs.append((c,evaluate(train_fn('zca',c),seeds=SEEDS)))
    best_cfg,best=max(idea_runs,key=lambda x:x[1]['mean']) if False else min(idea_runs,key=lambda x:x[1]['mean'])
    rep=make_report('vision','cnn_small',base,best[1],extra=signature(best_cfg))
    rep['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
    rep['protocol_notes']='Vision chosen because anchored whitening is a channel normalization mechanism; same CNN and training budget used for BN and ZCA.'
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()