import json, math, random, time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED = 1729

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


def gaussian_kernel(sigma, device, dtype, truncate=3.0):
    radius = max(1, int(math.ceil(truncate*sigma)))
    x = torch.arange(-radius, radius+1, device=device, dtype=dtype)
    g = torch.exp(-0.5*(x/sigma)**2); g = g/g.sum()
    k = g[:, None] * g[None, :]
    return k


def smooth_depthwise(x, sigma):
    b,c,h,w = x.shape
    k = gaussian_kernel(float(sigma), x.device, x.dtype)
    p = k.shape[0]//2
    weight = k[None,None].expand(c,1,-1,-1)
    return F.conv2d(F.pad(x, (p,p,p,p), mode='replicate'), weight, groups=c)


def besov_regularizer(x, sigma0=1.0, r=2.0, J=3, s=0.5, return_scales=False):
    scales = [smooth_depthwise(x, sigma0*(r**j)) for j in range(J+1)]
    ds = [scales[j]-scales[j+1] for j in range(J)]
    # RMS over spatial dimensions, then average channels and batch for a scalar.
    vals = []
    for j,d in enumerate(ds):
        rms = d.pow(2).mean(dim=(-2,-1)).sqrt() # B,C
        vals.append((2.0**(j*s))*rms)
    stacked = torch.stack(vals, dim=0)
    reg = torch.sqrt((stacked**2).sum(dim=0)+1e-12).mean()
    if return_scales: return reg, scales, ds, stacked
    return reg


def math_check():
    seed_all(); x = torch.randn(2,3,32,32)
    reg, scales, ds, weighted = besov_regularizer(x, return_scales=True)
    residual = (scales[0]-scales[-1]-sum(ds)).abs().max().item()
    # A sinusoid's response should move toward earlier (fine) scales as frequency rises.
    yy,xx = torch.meshgrid(torch.arange(32), torch.arange(32), indexing='ij')
    responses=[]
    for cycles in [1,4,10]:
        z = torch.sin(2*math.pi*cycles*xx/32).float()[None,None]
        _,_,_,w = besov_regularizer(z, return_scales=True)
        responses.append(w.detach().numpy().ravel().tolist())
    # Compare s=0 versus positive s on the same nonzero multi-scale signal.
    low = torch.sin(2*math.pi*2*xx/32).float()[None,None]
    high = torch.sin(2*math.pi*12*xx/32).float()[None,None]
    r0_low=besov_regularizer(low,s=0.0).item(); r0_high=besov_regularizer(high,s=0.0).item()
    rp_low=besov_regularizer(low,s=.5).item(); rp_high=besov_regularizer(high,s=.5).item()
    return {'telescoping_max_abs':residual, 'sinusoid_weighted_by_scale':responses,
            's0_low_high':[r0_low,r0_high], 's05_low_high':[rp_low,rp_high]}

class TinyNet(nn.Module):
    def __init__(self, use_reg=False, s=.5, lam=0.01):
        super().__init__(); self.use_reg=use_reg; self.s=s; self.lam=lam
        self.c1=nn.Conv2d(1,16,3,padding=1); self.c2=nn.Conv2d(16,32,3,padding=1)
        self.head=nn.Linear(32*8*8,2)
    def forward(self,x):
        x=F.relu(self.c1(x)); x=F.avg_pool2d(x,2)
        feat=F.relu(self.c2(x)); x=F.avg_pool2d(feat,2)
        logits=self.head(x.flatten(1))
        reg=besov_regularizer(feat,s=self.s) if self.use_reg else feat.new_zeros(())
        return logits,reg

def dataset(n, seed):
    g=torch.Generator().manual_seed(seed)
    # Two classes: a persistent broad bar versus a fine checker/stripe texture,
    # with nuisance translations and Gaussian noise.
    imgs=torch.randn(n,1,32,32,generator=g)*.12; ys=torch.randint(0,2,(n,),generator=g)
    for i,y in enumerate(ys.tolist()):
        shift=int(torch.randint(-2,3,(1,),generator=g));
        if y==0:
            imgs[i,:,9+shift:23+shift,14:18] += 0.9
            imgs[i,:,14:18,9+shift:23+shift] += 0.9
        else:
            yy,xx=torch.meshgrid(torch.arange(32),torch.arange(32),indexing='ij')
            imgs[i,0] += .38*torch.sin(2*math.pi*7*xx/32)
            imgs[i,0] += .38*torch.sin(2*math.pi*7*yy/32)
    return imgs,ys

def run_train(use_reg, lam, s, train, test, device):
    seed_all(99); model=TinyNet(use_reg,s,lam).to(device)
    opt=torch.optim.AdamW(model.parameters(),lr=2e-3,weight_decay=1e-4)
    x,y=train; xt,yt=test; x=x.to(device); y=y.to(device); xt=xt.to(device); yt=yt.to(device)
    losses=[]; accs=[]; regs=[]
    for step in range(160):
        idx=torch.randperm(len(x),device=device)[:64]
        logits,reg=model(x[idx]); loss=F.cross_entropy(logits,y[idx])+lam*reg
        opt.zero_grad(); loss.backward(); opt.step()
        if step in [39,79,159]:
            with torch.no_grad():
                z,r=model(xt); acc=(z.argmax(1)==yt).float().mean().item()
            losses.append(float(loss.item())); accs.append(acc); regs.append(float(reg.item()))
    return {'val_acc':accs,'train_loss':losses,'reg':regs}

def main():
    seed_all(); device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        train=dataset(512,11); test=dataset(256,12)
        # Keep data on CPU until training; CUDA failures are handled below.
        out={'device':str(device),'math':math_check()}
        try:
            base=run_train(False,0,.5,train,test,device)
            idea=run_train(True,.01,.5,train,test,device)
        except Exception as e:
            if device.type=='cuda':
                device=torch.device('cpu'); base=run_train(False,0,.5,train,test,device); idea=run_train(True,.01,.5,train,test,device)
                out['cuda_error']=repr(e)
            else: raise
        out.update({'baseline':base,'idea':idea})
        Path('results.json').write_text(json.dumps(out,indent=2))
        print(json.dumps(out,indent=2))
    except RuntimeError as e:
        if 'out of memory' in str(e).lower() and torch.cuda.is_available():
            torch.cuda.empty_cache(); print('CUDA OOM; rerun with CPU')
            raise
        raise
if __name__=='__main__': main()