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

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

def denominators(J, cutoff=None):
    # Golden-ratio conjugate: all continued-fraction partial quotients are 1.
    qs=[]; qm1, q=0, 1
    for _ in range(J):
        if cutoff is not None and q > cutoff: break
        qs.append(q)
        qn = q + qm1
        qm1, q = q, qn
    return np.asarray(qs, dtype=np.int64)

def lacunary_features(x, qs, square_root=True):
    x = np.asarray(x, dtype=np.float64).reshape(-1, 1)
    q = np.asarray(qs, dtype=np.float64).reshape(1, -1)
    scale = 1/np.sqrt(q) if square_root else 1/q
    z = 2*np.pi*x*q
    return np.concatenate([np.cos(z)*scale, np.sin(z)*scale], axis=1)

def h_signal(x, qs):
    q=np.asarray(qs,dtype=float)
    return (np.cos(2*np.pi*np.asarray(x).reshape(-1,1)*q)/q).sum(1)

def regularity_check():
    # A direct finite approximation: dyadic modulus, and growth of the
    # finite-truncation Lipschitz upper bound sum_j 2*pi (each term contributes 2*pi).
    qs=denominators(25)
    grid=np.linspace(0,1,200001)
    vals=h_signal(grid,qs)
    deltas=2.0**np.arange(-12,-3,dtype=float)
    mod=[]
    for d in deltas:
        k=max(1, int(round(d*len(grid))))
        mod.append(float(np.max(np.abs(vals[k:]-vals[:-k]))))
    mod=np.asarray(mod)
    slope=float(np.polyfit(np.log(deltas), np.log(mod), 1)[0])
    # For the truncated series, derivative absolute sum is exactly 2*pi*J.
    growth=[]
    for J in [4,8,12,16,20,25]: growth.append((J, 2*np.pi*J))
    # Verify recursion and the claimed absolutely summable coefficient series.
    recursion_ok=bool(np.all(qs[2:] == qs[1:-1]+qs[:-2]))
    return {'q_first_12': qs[:12].tolist(), 'recursion_ok':recursion_ok,
            'sum_inverse_q':float(np.sum(1/qs)), 'dyadic_deltas':deltas.tolist(),
            'modulus':mod.tolist(), 'loglog_slope':slope,
            'finite_lipschitz_bound_2piJ':growth}

class MLP(nn.Module):
    def __init__(self,d):
        super().__init__(); self.net=nn.Sequential(nn.Linear(d,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh(),nn.Linear(64,1))
    def forward(self,x): return self.net(x)

def train_model(kind, device):
    # Target is exactly a longer continued-fraction series; held-out region tests
    # whether the structured features extrapolate rather than merely interpolate.
    qtarget=denominators(10)
    xtr=np.linspace(0, .70, 448, endpoint=False)
    xte=np.linspace(.70, 1.0, 192, endpoint=False)
    ytr=h_signal(xtr,qtarget); yte=h_signal(xte,qtarget)
    qfeat=denominators(8) # [1,1,2,3,5,8,13,21], 16 channels for all encoded models
    def make(x):
        if kind=='raw': return x[:,None]
        if kind=='nerf':
            # 8 standard dyadic frequencies, same 16 sinusoidal channels.
            z=2*np.pi*x[:,None]*(2.0**np.arange(8)[None,:])
            return np.concatenate([np.sin(z),np.cos(z)],1)
        return np.concatenate([x[:,None],lacunary_features(x,qfeat,True)],1)
    Xtr=torch.tensor(make(xtr),dtype=torch.float32,device=device); Ytr=torch.tensor(ytr[:,None],dtype=torch.float32,device=device)
    Xte=torch.tensor(make(xte),dtype=torch.float32,device=device); Yte=torch.tensor(yte[:,None],dtype=torch.float32,device=device)
    torch.manual_seed(SEED); model=MLP(Xtr.shape[1]).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    for step in range(1800):
        opt.zero_grad(); loss=((model(Xtr)-Ytr)**2).mean(); loss.backward(); opt.step()
    with torch.no_grad():
        tr=float(((model(Xtr)-Ytr)**2).mean().sqrt().cpu()); te=float(((model(Xte)-Yte)**2).mean().sqrt().cpu())
    return {'train_rmse':tr,'extrapolation_rmse':te,'input_dim':int(Xtr.shape[1]),'parameters':sum(p.numel() for p in model.parameters())}

def main():
    try:
        device='cuda' if torch.cuda.is_available() else 'cpu'
        # CUDA can fail on a shared device; this makes the experiment robust.
        out={'device':device,'math':regularity_check()}
        try: out['benchmark']={k:train_model(k,device) for k in ['raw','nerf','lacunary']}
        except Exception as e:
            device='cpu'; out['device_fallback']=str(e); out['benchmark']={k:train_model(k,device) for k in ['raw','nerf','lacunary']}
    except Exception as e:
        raise
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()