import json, math, random
from pathlib import Path
import numpy as np

SEED = 2096
np.random.seed(SEED); random.seed(SEED)
OUT = Path('results.json')

def toy_scaling():
    # g(t)=a+b*t+c*t^2 gives exact linear-extrapolation residual 2|c|h^2.
    hs = np.logspace(-3, -0.3, 12)
    a, b, c = 0.7, -0.4, 1.3
    vals = []
    for h in hs:
        old2 = a + b*(-h) + c*(-h)**2
        old1 = a
        current = a + b*h + c*h*h
        predicted = old1 + (old1-old2)
        vals.append(abs(current-predicted))
    slope = float(np.polyfit(np.log(hs), np.log(vals), 1)[0])
    C, eps = 2.0*abs(c), 0.05
    claimed_target = math.sqrt(eps/C)
    # This is exactly the formula in the idea: h_next=sqrt(eps/||r||),
    # with projection and ratio limiting. Its fixed point is fourth-root,
    # because ||r||=C*h^2; this is tested separately from the claimed target.
    h, hmin, hmax, q, delta = 0.8, 1e-4, 2.0, 1.35, 1e-12
    trace=[]
    for _ in range(100):
        residual=C*h*h
        raw=math.sqrt(eps/(residual+delta))
        nh=float(np.clip(raw,hmin,hmax)); nh=float(np.clip(nh,h/q,q*h))
        trace.append((h,residual,nh)); h=nh
    observed_h=float(np.mean([z[0] for z in trace[-10:]]))
    observed_r=float(np.mean([z[1] for z in trace[-10:]]))
    formula_fixed_point=(eps/C)**0.25
    rng=np.random.default_rng(SEED); h=.2; noisy_ratios=[]
    for _ in range(100):
        measured=max(1e-9,C*h*h*(1+.15*rng.normal()))
        nh=float(np.clip(math.sqrt(eps/measured),hmin,hmax))
        nh=float(np.clip(nh,h/q,q*h)); noisy_ratios.append(nh/h); h=nh
    allrat=np.array([z[2]/z[0] for z in trace]+noisy_ratios)
    return {'residual_scaling': {'observed_slope':slope,'predicted_slope':2.0,'absolute_error':abs(slope-2)},
      'scheduler_target': {'C':C,'epsilon':eps,'claimed_sqrt_target':claimed_target,
        'observed_h':observed_h,'observed_residual':observed_r,
        'formula_fixed_point_fourth_root':formula_fixed_point,
        'observed_vs_fourth_root_relative_error':abs(observed_h-formula_fixed_point)/formula_fixed_point,
        'claimed_target_relative_error':abs(observed_h-claimed_target)/claimed_target,
        'note':'The implemented formula lacks the multiplicative h_k present in the paper-style ratio controller.'},
      'ratio_sweep': {'observed_min':float(allrat.min()),'observed_max':float(allrat.max()),'predicted_bounds':[1/q,q]}}

def benchmark():
    import torch
    torch.manual_seed(SEED); np.random.seed(SEED)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        n=1200; rng=np.random.default_rng(SEED)
        x=np.concatenate([rng.normal([-1.,-1.],.75,(n//2,2)),rng.normal([1.,1.],.75,(n//2,2))]).astype('float32')
        y=np.concatenate([np.zeros(n//2),np.ones(n//2)]).astype('int64'); p=rng.permutation(n); x=x[p]; y=y[p]
        X=torch.tensor(x,device=device); Y=torch.tensor(y,device=device)
        def run(kind):
            torch.manual_seed(SEED+3)
            model=torch.nn.Sequential(torch.nn.Linear(2,24),torch.nn.Tanh(),torch.nn.Linear(24,2)).to(device)
            lossfn=torch.nn.CrossEntropyLoss(); params=list(model.parameters()); prev=[]; intervals=[.08,.08]
            h=.08; eps=.003; hmin=.015; hmax=.20; q=1.25; beta=.8; R=None; losses=[]; hs=[]; accs=[]
            for k in range(260):
                ix=torch.randint(0,n,(64,),device=device); model.zero_grad(set_to_none=True)
                loss=lossfn(model(X[ix]),Y[ix]); loss.backward(); g=torch.cat([p.grad.detach().reshape(-1) for p in params])
                if kind=='adaptive' and len(prev)>=2:
                    pred=prev[-1]+(h/intervals[-1])*(prev[-1]-prev[-2]); res=torch.linalg.vector_norm(g-pred).item()
                    R=res if R is None else beta*R+(1-beta)*res
                    nh=float(np.clip(math.sqrt(eps/(R+1e-8)),hmin,hmax)); nh=float(np.clip(nh,h/q,q*h))
                else: nh=h
                with torch.no_grad():
                    for z in params: z -= nh*z.grad
                prev.append(g); prev=prev[-3:]; intervals.append(nh); intervals=intervals[-3:]
                h=nh if kind=='adaptive' else (.08*(.5*(1+math.cos(math.pi*(k+1)/260)) if kind=='cosine' else 1.))
                losses.append(float(loss)); hs.append(nh)
                with torch.no_grad(): accs.append(float((model(X).argmax(1)==Y).float().mean()))
            return {'final_loss':float(np.mean(losses[-20:])),'final_accuracy':float(np.mean(accs[-20:])),'best_loss':float(min(losses)),'mean_interval':float(np.mean(hs)),'interval_min':float(min(hs)),'interval_max':float(max(hs))}
        out={k:run(k) for k in ['fixed','cosine','adaptive']}; out['device']=device; return out
    except Exception as e:
        if device=='cuda':
            torch.cuda.empty_cache(); return {'cuda_error':str(e),'fallback':'run with CUDA_VISIBLE_DEVICES=-1'}
        return {'error':str(e)}

def main():
    data={'seed':SEED,'mechanism':toy_scaling()}
    try: data['benchmark']=benchmark()
    except Exception as e: data['benchmark_error']=repr(e)
    OUT.write_text(json.dumps(data,indent=2)); print(json.dumps(data,indent=2))
if __name__=='__main__': main()