import math, json, random
import numpy as np
import torch

SEED=417
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def f(r): return 1.0/((r+2.0)*math.log(r+2.0)**2)

def math_check():
    # Add the analytic integral remainder after the finite summation cutoff.
    M=2_000_000
    rs=np.arange(1,M+1,dtype=np.float64)
    vals=1.0/((rs+2.0)*np.log(rs+2.0)**2)
    tails=[]
    for R in [8,16,32,64,128,256,1024]:
        finite=float(vals[R:].sum())
        remainder=1.0/math.log(M+2.0)
        corrected=finite+remainder
        estimate=1.0/math.log(R+2.0)
        tails.append({'R':R,'tail_with_remainder':corrected,'integral_estimate':estimate,'ratio':corrected/estimate})
    bands=[]
    for k in range(1,9):
        lo,hi=2**k,2**(k+1)
        exact=sum(f(r) for r in range(lo,hi))
        integral=1.0/math.log(lo+2.0)-1.0/math.log(hi+2.0)
        asymptotic=1.0/(k*(k+1)*math.log(2.0))
        bands.append({'k':k,'exact':exact,'shifted_integral':integral,'asymptotic_dyadic_mass':asymptotic,'exact_over_integral':exact/integral})
    return {'tails':tails,'bands':bands}

def exact_kernel(L):
    # symmetric normalized off-diagonal kernel, including only distances available in length L
    a=np.array([0.0]+[f(r) for r in range(1,L)],dtype=np.float64)
    return a/(2*a[1:].sum())

def approx_dyadic(L):
    # Effective coefficient for each distance induced by band averaging and fixed band masses.
    q=np.zeros(L)
    for k in range(int(math.log2(L))):
        lo,hi=2**k,min(2**(k+1),L)
        mass=1.0/((k+1)*(k+2)*math.log(2.0))
        q[lo:hi]+=mass/(hi-lo)
    q=q/(2*q[1:].sum())
    return q

def local(L,R=8):
    q=np.zeros(L); q[1:R+1]=1.0/(2*R); return q

def compare():
    rows=[]
    for L in [64,256,1024,4096]:
        p=exact_kernel(L)
        d=approx_dyadic(L); c=local(L)
        # relative L1 discrepancy and mass beyond the local radius
        rows.append({'L':L,'dyadic_l1_error':float(np.abs(p-d).sum()),'local_l1_error':float(np.abs(p-c).sum()),'true_mass_r_gt_8':float(p[9:].sum()),'dyadic_mass_r_gt_8':float(d[9:].sum())})
    return rows

print(json.dumps({'math':math_check(),'kernel_approximation':compare()},indent=2))