import os, sys, json, time, 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

SEED=115
SEEDS=tuple(range(8))
# The idea's binary tree is represented explicitly; K=8 keeps this bench small.
def tree_splits(k=8, lo=0., hi=1.):
    if k <= 1: return [0]
    a=(lo+hi)/2
    l=k//2; r=k-l
    return [1+x for x in tree_splits(l,lo,a)] + [1+x for x in tree_splits(r,a,hi)]

def math_check():
    ratios=np.array([abs(.5**(j+1)) / abs(.5**j) for j in range(1,8)])
    paths=tree_splits(8)
    assert np.allclose(ratios,.5) and max(paths)==3
    return {'geometric_ratio':float(ratios.max()),'expected_ratio':.5,
            'ratio_allclose':True,'K':8,'M':2,'path_split_counts':paths,
            'max_splits':max(paths),'N':3,'bound_holds':max(paths)<=3}

class DirectionalBlock(nn.Module):
    def __init__(self, c, mode, k=8, beam=2):
        super().__init__(); self.mode=mode; self.k=k; self.beam=beam
        self.slopes=torch.linspace(0,1,k)
        self.proj=nn.Conv2d(c,c,1,bias=False)
        self.route=nn.Conv2d(c, k, 1)
        self.last_samples=0; self.last_selected=0; self.last_entropy=0.
    def forward(self,x):
        b,c,h,w=x.shape; z=self.proj(x)
        # One shared local operator samples every angular leaf; routing only
        # changes which leaf responses enter the output.
        responses=[]
        yy=torch.linspace(-1,1,h,device=x.device); xx=torch.linspace(-1,1,w,device=x.device)
        gy,gx=torch.meshgrid(yy,xx,indexing='ij'); base=torch.stack((gx,gy),-1)[None].expand(b,-1,-1,-1)
        for slope in self.slopes.to(x.device):
            # short line: offsets -1,0,+1, with horizontal displacement
            acc=0.
            for t in (-1.,0.,1.):
                grid=base.clone(); grid[...,0]=grid[...,0]+t*2/(w-1); grid[...,1]=grid[...,1]+t*float(slope)*2/(h-1)
                acc=acc+F.grid_sample(z,grid,mode='bilinear',padding_mode='border',align_corners=True)
            responses.append(acc/3.)
        r=torch.stack(responses,1) # B,K,C,H,W
        logits=self.route(x)
        p=logits.softmax(1)
        if self.mode=='dense':
            out=(r*p[:, :, None]).sum(1); chosen=self.k
        else:
            vals,idx=torch.topk(p,self.beam,dim=1)
            mask=torch.zeros_like(p).scatter(1,idx,1.)
            out=(r*(p*mask)[:, :, None]).sum(1)/(vals.sum(1)[:,None]+1e-6); chosen=self.beam
        self.last_samples=int(b*h*w*chosen*3); self.last_selected=chosen
        self.last_entropy=float((-(p*(p+1e-8).log()).sum(1)).mean().detach().cpu())
        return F.relu(out+x)

class DirectionalCNN(nn.Module):
    def __init__(self, out_dim, mode):
        super().__init__(); self.mode=mode
        self.stem=nn.Sequential(nn.Conv2d(3,16,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),
                                nn.Conv2d(16,32,3,padding=1),nn.ReLU(),nn.MaxPool2d(2))
        self.dir=DirectionalBlock(32,mode)
        self.tail=nn.Sequential(nn.Conv2d(32,64,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(64*4*4,64),nn.ReLU(),nn.Linear(64,out_dim))
    def forward(self,x): return self.tail(self.dir(self.stem(x)))

def train_one(seed, mode, lr, epochs=5, return_model=False):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    d=get_dataset('vision',seed,n_train=400,n_test=200)
    net=DirectionalCNN(d['out_dim'],mode)
    net,metric,hist=train_model(net,d,epochs=epochs,lr=lr,batch=64,log=lambda *_:None)
    if return_model: return metric,net
    return metric

def main():
    check=math_check(); print('math_check',json.dumps(check))
    # Union parity: baseline evaluates every lr offered to idea.
    grid=[{'lr':1e-3},{'lr':3e-3},{'lr':5e-3}]
    base=sweep_baseline(lambda cfg: (lambda s: train_one(s,'dense',cfg['lr'])),grid)
    # Idea is run at best baseline lr and two nearby values (all in union).
    idea_cfgs=[c['cfg'] for c in base['sweep']]
    idea_runs=[]
    for cfg in idea_cfgs:
        vals=[train_one(s,'routed',cfg['lr']) for s in SEEDS]
        idea_runs.append({'cfg':cfg,'result':{'mean':float(np.mean(vals)),'std':float(np.std(vals)), 'per_seed':vals,'n':len(vals)}})
    best=min(idea_runs,key=lambda z:z['result']['mean']); idea=best['result']
    # Signature is measured on trained systems, not an analytic toy identity.
    sm=[]
    for s in SEEDS[:2]:
        bm,bnet=train_one(s,'dense',base['best_cfg']['lr'],return_model=True)
        im,inet=train_one(s,'routed',base['best_cfg']['lr'],return_model=True)
        sm.append({'seed':s,'dense_metric':bm,'routed_metric':im,
                   'dense_samples_per_query':bnet.dir.last_samples/(400*8*8),
                   'routed_samples_per_query':inet.dir.last_samples/(400*8*8),
                   'observed_reduction':bnet.dir.last_samples/max(1,inet.dir.last_samples),
                   'routing_entropy':inet.dir.last_entropy})
    sig={'prediction':'top-2 of K=8 uses 4x fewer directional samples','K':8,'beam':2,
         'predicted_reduction':4.0,'observed':sm,'confirmed':all(abs(x['observed_reduction']-4)<1e-6 for x in sm)}
    rep=make_report('vision','directional_cnn',base,idea,{'math_check':check,'idea_sweep':idea_runs,
        'mechanism_signature':sig,'custom_track':None,'structural_match':'spatial CIFAR images and convolutional directional fields'})
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))
if __name__=='__main__': main()