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

SEED=417
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda': torch.cuda.empty_cache()
except Exception:
    device=torch.device('cpu')

def kval(r, r0=2.0):
    return 1.0/((r+r0)*math.log(r+r0)**2)

def math_check():
    rs=np.arange(1,300001,dtype=np.float64)
    kv=1/((rs+2)*np.log(rs+2)**2)
    tail=[]
    for R in [8,16,32,64,128,256]:
        s=float(kv[R:].sum()); q=1/math.log(R+2)
        tail.append([R,s,q,s/q])
    bands=[]
    for k in range(7):
        lo,hi=2**k,2**(k+1)
        exact=float(kv[lo-1:hi-1].sum())
        formula=1/((k+1)*(k+2)*math.log(2))
        bands.append([k,exact,formula,exact/formula])
    return {'tail_R_discrete_integral_ratio':tail,'bands_k_exact_formula_ratio':bands}

def make_weights(L):
    w=torch.zeros(L,device=device)
    for r in range(1,L): w[r]=kval(r)
    return w/(2*w.sum())

def long_target(x,w):
    # x [B,L,1], symmetric zero-padded convolution
    B,L,D=x.shape; y=torch.zeros_like(x)
    for r in range(1,L):
        y[:,r:]+=w[r]*x[:,:L-r]; y[:,:L-r]+=w[r]*x[:,r:]
    return y

def dyadic_mix(x):
    B,L,D=x.shape; out=torch.zeros_like(x); pref=torch.cat([torch.zeros(B,1,D,device=x.device),x.cumsum(1)],1)
    K=int(math.log2(L))
    for k in range(K):
        lo,hi=2**k, min(2**(k+1),L)
        # clipped ranges, with exact valid-count normalization per position
        idx=torch.arange(L,device=x.device)
        a=(idx-hi).clamp(min=0); b=(idx-lo).clamp(min=0)
        left=(pref[:,b+1]-pref[:,a+1]) # positions [i-hi+1,i-lo], then remove invalid below zero
        left_count=(b-a).clamp(min=1).float()[None,:,None]
        # Correct range is j in [i-hi, i-lo], exclusive invalid j<0.
        left=(pref[:,(idx-lo+1).clamp(min=0)]-pref[:,(idx-hi).clamp(min=0)])/left_count
        a=(idx+lo).clamp(max=L); b=(idx+hi).clamp(max=L)
        right=(pref[:,b]-pref[:,a])/(b-a).clamp(min=1).float()[None,:,None]
        ak=1/((k+1)*(k+2)*math.log(2))
        out += ak*(left+right)
    return out

def local_mix(x,r=8):
    B,L,D=x.shape; out=torch.zeros_like(x)
    for d in range(1,r+1):
        out[:,d:]+=x[:,:L-d]; out[:,:L-d]+=x[:,d:]
    return out/(2*r)

class Dyadic(nn.Module):
    def __init__(self): super().__init__(); self.scale=nn.Parameter(torch.tensor(.5))
    def forward(self,x): return x+self.scale*dyadic_mix(x)
class Local(nn.Module):
    def __init__(self,r=8): super().__init__(); self.scale=nn.Parameter(torch.tensor(.5)); self.r=r
    def forward(self,x): return x+self.scale*local_mix(x,self.r)

def train(model, L, steps=180, batch=64):
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=.03); w=make_weights(L)
    losses=[]; t=time.time()
    for s in range(steps):
        x=torch.randn(batch,L,1,device=device); y=long_target(x,w)
        pred=model(x); loss=F.mse_loss(pred,y); opt.zero_grad(); loss.backward(); opt.step()
        losses.append(float(loss))
    with torch.no_grad():
        x=torch.randn(512,L,1,device=device); y=long_target(x,w); val=float(F.mse_loss(model(x),y))
    return {'final_train':losses[-1],'val_mse':val,'seconds':time.time()-t,'learned_scale':float(model.scale.detach().cpu())}

def main():
    out={'device':str(device),'math':math_check(),'experiments':{}}
    for L in [64,256]:
        # same architecture/optimization; local control has same single scalar parameter
        out['experiments'][str(L)]={'local_radius8':train(Local(8),L),'dyadic_tail':train(Dyadic(),L)}
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()