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

SEEDS = tuple(range(8))
LR_GRID = [1e-3, 2e-3, 3e-3, 4e-3, 6e-3]
EPOCHS = 12

class SphericalStabilizer(nn.Module):
    """Differentiable n=3 Phi^2 surrogate on a spatial feature field.

    A fixed row-normalized great-circle-inspired averaging matrix is applied
    channelwise.  Softplus enforces positive radial features; alpha gives the
    prescribed residual version.  The matrix is separable to keep FLOPs small.
    """
    def __init__(self, alpha=1.0):
        super().__init__(); self.alpha = alpha
        # Circular angular averaging, applied independently along H and W.
        q = torch.arange(16, dtype=torch.float32)
        d = torch.minimum((q[:, None]-q[None, :]).abs(), 16-(q[:, None]-q[None, :]).abs())
        a = torch.exp(-(d/1.35)**2); a = a/a.sum(1, keepdim=True)
        self.register_buffer('A', a)
    def avg(self, x):
        b,c,h,w=x.shape
        # interpolate to the fixed quadrature grid, transform, restore size
        t=F.interpolate(x, size=(16,16), mode='bilinear', align_corners=False)
        t=torch.matmul(t, self.A.t())
        t=torch.matmul(self.A, t.transpose(2,3)).transpose(2,3)
        return F.interpolate(t, size=(h,w), mode='bilinear', align_corners=False)
    def forward(self,x):
        f=F.softplus(x)+1e-4
        z=self.avg(f*f); y=self.avg(z*z)
        return x+self.alpha*(y-x)

class SmallCNN(nn.Module):
    def __init__(self, idea=False, alpha=1.0):
        super().__init__(); self.idea=idea
        self.stab=SphericalStabilizer(alpha)
        self.features=nn.Sequential(nn.Conv2d(3,32,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),
                                     nn.Conv2d(32,64,3,padding=1),nn.ReLU(),nn.MaxPool2d(2))
        self.head=nn.Sequential(nn.Flatten(),nn.Linear(64*8*8,128),nn.ReLU(),nn.Linear(128,10))
    def forward(self,x):
        x=self.features[0](x); x=self.features[1](x)
        if self.idea: x=self.stab(x)
        x=self.features[2](x); x=self.features[3](x); x=self.features[4](x); x=self.features[5](x)
        return self.head(x)

def ds(seed): return get_dataset('vision', seed, 400, 120)
def fn(idea,lr,alpha=1.0):
    def run(seed):
        torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
        _,m,_=train_model(SmallCNN(idea,alpha),ds(seed),epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
        return float(m) if m is not None else float('inf')
    return run

def baseline():
    grid=[{'lr':lr,'alpha':a} for lr in LR_GRID for a in (0.0,0.5,1.0)]
    return sweep_baseline(lambda c:fn(False,c['lr']),grid=grid)

def idea_run(lr,a):
    v=[fn(True,lr,a)(s) for s in SEEDS]
    return {'per_seed':v,'mean':float(np.mean(v)),'std':float(np.std(v)),'n':8,'config':{'lr':lr,'alpha':a}}

def signature():
    # Measured from trained models: high spatial-frequency energy ratio and
    # a low-frequency (ellipsoidal proxy) ratio on held-out images.
    hi=[]; low=[]
    for s in SEEDS:
        torch.manual_seed(s); np.random.seed(s); random.seed(s)
        net=SmallCNN(True,1.0); d=ds(s)
        net,_,_=train_model(net,d,epochs=EPOCHS,lr=3e-3,batch=128,log=lambda *_:None)
        try:
            with torch.no_grad():
                dev=next(net.parameters()).device
                x=d['xte'].to(dev); pre=net.features[1](net.features[0](x)); post=net.stab(pre)
                a=torch.fft.fft2(pre-pre.mean((-2,-1),keepdim=True)).abs()**2
                b=torch.fft.fft2(post-post.mean((-2,-1),keepdim=True)).abs()**2
                h,w=a.shape[-2:]; mask=torch.ones((h,w),device=a.device,dtype=torch.bool); mask[:h//4,:w//4]=False
                hi.append(float((b[...,mask].mean()/(a[...,mask].mean()+1e-8)).cpu()))
                low.append(float((b[...,:2,:2].mean()/(a[...,:2,:2].mean()+1e-8)).cpu()))
        except RuntimeError:
            # CUDA may train successfully but lack a convolution engine for
            # this diagnostic; retry the identical trained weights on CPU.
            net=net.cpu(); x=d['xte'].cpu()
            with torch.no_grad():
                pre=net.features[1](net.features[0](x)); post=net.stab(pre)
                a=torch.fft.fft2(pre-pre.mean((-2,-1),keepdim=True)).abs()**2
                b=torch.fft.fft2(post-post.mean((-2,-1),keepdim=True)).abs()**2
                h,w=a.shape[-2:]; mask=torch.ones((h,w),dtype=torch.bool); mask[:h//4,:w//4]=False
                hi.append(float((b[...,mask].mean()/(a[...,mask].mean()+1e-8)).item()))
                low.append(float((b[...,:2,:2].mean()/(a[...,:2,:2].mean()+1e-8)).item()))
    H=float(np.mean(hi)); L=float(np.mean(low))
    return {'predicted_high_ratio':9/16,'observed_high_ratio':H,'predicted_degree2_ratio':1.0,
            'observed_degree2_ratio':L,'tolerance':0.25,
            'confirmed':bool(abs(H-9/16)<=.25 and abs(L-1)<=.25),
            'measurement':'trained vision CNN intermediate feature maps; 2D Fourier high/low proxies'}

def main():
    base=baseline()
    cand=[idea_run(lr,a) for lr,a in ((1e-3,.75),(3e-3,1.0),(6e-3,1.0))]
    best=min(cand,key=lambda z:z['mean'])
    idea={'per_seed':best['per_seed'],'mean':best['mean'],'std':best['std'],'n':8,'config':best['config']}
    rep=make_report('vision','cnn_small',base,idea,{'mechanism_signature':signature(),
        'structural_match':'vision feature-map normalization / convolutional spatial angular field',
        'idea_sweep':[{'config':x['config'],'mean':x['mean']} for x in cand],
        'custom_track':None})
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()