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

SEED0 = 1723
LRS = [1e-3, 3e-3, 1e-2]
EPOCHS = 5
NTR, NTE = 400, 200

class TernaryMeanPool(nn.Module):
    """Standard pooling baseline used by the bench CNN."""
    def forward(self, x):
        return F.max_pool2d(x, 2)

class FractalPool(nn.Module):
    """Two-channel ternary renormalization pooling, with bounded defect route."""
    def __init__(self, channels, width=16):
        super().__init__()
        self.x_raw = nn.Parameter(torch.tensor(-1.0))
        self.lam_raw = nn.Parameter(torch.tensor(-0.15))
        self.defect = nn.Sequential(nn.Conv2d(3*channels, width, 1), nn.GELU(),
                                    nn.Conv2d(width, channels, 1), nn.Tanh())
        self.gate = nn.Sequential(nn.Conv2d(3*channels, channels, 1), nn.Sigmoid())
        self.last_rho = 0.0
    def forward(self, x):
        if x.shape[-1] % 2 or x.shape[-2] % 2:
            x = F.pad(x, (0, x.shape[-1] % 2, 0, x.shape[-2] % 2))
        a = x[..., 0::2, 0::2]; b = x[..., 0::2, 1::2]; c = x[..., 1::2, 0::2]
        # neutral and defect channels: disagreement is a child-minus-neutral signal.
        neutral0 = (a + b + c) / 3.0
        da, db, dc = a-neutral0, b-neutral0, c-neutral0
        lam, xx = F.softplus(self.lam_raw), F.softplus(self.x_raw)
        neutral = lam.pow(3)*a*b*c + 2*xx.pow(3)*da*db*dc
        neutral = neutral / (neutral.square().mean(1, keepdim=True).add(1e-5).sqrt())
        defect_in = torch.cat((da, db, dc), 1)
        proposed = self.defect(defect_in)
        gate = self.gate(defect_in)
        defect = gate*proposed + (1-gate)*neutral0
        self._last_neutral, self._last_defect = neutral, defect
        return neutral + defect
    def coeffs(self):
        return float(F.softplus(self.x_raw)), float(F.softplus(self.lam_raw))

class CNN(nn.Module):
    def __init__(self, idea=False):
        super().__init__(); P = FractalPool if idea else TernaryMeanPool
        self.c1=nn.Conv2d(3,32,3,padding=1); self.c2=nn.Conv2d(32,64,3,padding=1); self.c3=nn.Conv2d(64,96,3,padding=1)
        self.p1=P(32) if idea else P(); self.p2=P(64) if idea else P(); self.p3=P(96) if idea else P()
        self.fc=nn.Sequential(nn.Flatten(),nn.Linear(96*4*4,128),nn.ReLU(),nn.Linear(128,10))
    def forward(self,x):
        x=F.relu(self.c1(x)); x=self.p1(x); x=F.relu(self.c2(x)); x=self.p2(x); x=F.relu(self.c3(x)); x=self.p3(x); return self.fc(x)

def math_check():
    a,b=0.43,0.27; lam,x=0.8,0.35
    f=lambda u: lam**3*u[0]**3+2*x**3*u[1]**3
    eps=1e-5; fd=np.array([(f((a+eps,b))-f((a-eps,b)))/(2*eps),(f((a,b+eps))-f((a,b-eps)))/(2*eps)])
    exact=np.array([3*lam**3*a*a,6*x**3*b*b])
    return {'analytic_jacobian':exact.tolist(),'finite_difference':fd.tolist(),'max_abs_error':float(np.max(np.abs(exact-fd))), 'cubic_gain_ratio':float(exact[0]/(3*a*a))}

def run_cfg(idea, lr, seed, keep=False):
    torch.manual_seed(10000+seed); np.random.seed(10000+seed); random.seed(10000+seed)
    d=get_dataset('vision',seed,n_train=NTR,n_test=NTE); m=CNN(idea=idea)
    net,metric,hist=train_model(m,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
    return (float(metric),net,d) if keep else float(metric)

def full(idea,lr):
    vals=[]; models=[]
    for s in range(8):
        r=run_cfg(idea,lr,s,keep=idea)
        if idea: vals.append(r[0]); models.append(r)
        else: vals.append(r)
    return {'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)}, models

def signature(models):
    ratios=[]; rhos=[]
    for metric,net,d in models:
        net = net.to('cpu'); net.eval(); x=d['xte'][:8].to('cpu').clone(); noise=torch.randn_like(x)*1e-3
        with torch.no_grad():
            # Re-test trained network behavior, not an analytic toy: first pooling perturbation.
            z=F.relu(net.c1(x)); zn=F.relu(net.c1(x+noise));
            y=net.p1(z); yn=net.p1(zn)
            ratios.append(float((yn-y).norm()/(noise.norm()+1e-12)))
        rhos.append(net.p1.coeffs() if hasattr(net.p1,'coeffs') else None)
    vals=np.asarray(ratios)
    return {'prediction':'local Jacobian controls perturbation amplification; cubic neutral derivative scales as lambda^3',
            'observed_trained_pool_perturbation_gain_mean':float(vals.mean()),
            'observed_trained_pool_perturbation_gain_std':float(vals.std()),
            'trained_coefficients':rhos, 'predicted_lambda_cubic_gain':True,
            'confirmed': False}

def main():
    check=math_check(); print('MATH_CHECK',json.dumps(check))
    # Baseline sweep uses exactly the union of idea learning rates; its full result is best config.
    base=sweep_baseline(lambda cfg: (lambda s: run_cfg(False,cfg['lr'],s)), [{'lr':x} for x in LRS], seeds=(0,1,2,3))
    base_full_by_lr={str(lr):full(False,lr)[0] for lr in LRS}
    idea_candidates=[]
    for lr in LRS:
        r,_=full(True,lr); idea_candidates.append((r['mean'],lr,r))
    _,best_lr,idea_res=min(idea_candidates,key=lambda z:z[0]); _,models=full(True,best_lr)
    rep=make_report('vision','cnn_small',base,idea_res,{'math_check':check,'trained_behavior':signature(models),'baseline_full_by_lr':base_full_by_lr,'idea_sweep':[{'lr':lr,'mean':r['mean']} for _,lr,r in idea_candidates]})
    rep['protocol_note']='Vision selected because the idea changes convolutional hierarchical feature propagation; baseline and idea share convolution/classifier stages and differ only in pooling mechanism.'
    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
if __name__=='__main__': main()