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

SEED = 1296
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if device == "cuda": torch.cuda.empty_cache()
except Exception:
    device = "cpu"

def controller_step(theta, loss_fn, eta_prev, alpha=1.0, rho=.9, c=.1,
                    beta=.5, eta_min=1e-5, eta_max=2.0, eps=1e-12, max_bt=12):
    theta = theta.detach().requires_grad_(True)
    f0 = loss_fn(theta)
    g, = torch.autograd.grad(f0, theta)
    G = torch.linalg.vector_norm(g).item()
    if G < 1e-14:
        return theta.detach(), float(f0), 0., 0, G, 0.
    # Probe the currently maintained rate, then infer directional curvature.
    trial = theta - eta_prev * g
    ftrial = loss_fn(trial).detach()
    s_norm = eta_prev * G
    r = (ftrial.item() - f0.item() + eta_prev * G * G)
    Lhat = (1 + alpha) * max(r, 0.) / (s_norm ** (1 + alpha) + eps)
    eta_star = (G ** (1-alpha) / (Lhat + eps)) ** (1/alpha)
    eta = min(eta_max, max(eta_min, rho * eta_star))
    accepted = False
    fc = float('inf')
    for j in range(max_bt + 1):
        cand = theta - eta * g
        fc = loss_fn(cand).detach().item()
        if fc <= f0.item() - c * eta * G * G:
            accepted = True
            break
        eta *= beta
    if not accepted:
        eta = 0.
        cand = theta
    # Mild growth damping: preserve the accepted scale rather than blindly probing it.
    next_eta = min(eta_max, eta / .9) if accepted else max(eta_min, eta_prev * beta)
    return cand.detach(), fc if accepted else f0.item(), next_eta, int(accepted), G, Lhat

def quadratic_sweep():
    rows=[]
    # Prediction 1: alpha=1 Lhat = directional curvature q/G^2, independent of probe eta.
    for lam in [0.1, 0.3, 1., 3., 10.]:
        H=torch.diag(torch.tensor([lam, 2.*lam, .5*lam], dtype=torch.float64))
        theta=torch.tensor([1.2,-.7,.4], dtype=torch.float64)
        def loss(x): return .5*x @ H @ x
        g=H@theta; G=torch.linalg.vector_norm(g).item(); q=(g @ H @ g).item()
        exact=q/(G*G)
        vals=[]
        for probe in [.03,.2,1.1]:
            trial=theta-probe*g
            r=(loss(trial)-loss(theta)+probe*G*G).item()
            lh=2*max(r,0)/(probe*G)**2
            vals.append(lh)
        rows.append({'lambda_scale':lam,'exact_directional_curvature':exact,
                     'observed_Lhat_mean':float(np.mean(vals)), 'relative_error':abs(np.mean(vals)-exact)/exact,
                     'eta_star_prediction':1/exact, 'observed_eta_star':1/float(np.mean(vals))})
    # Prediction 2: accepted eta * directional curvature = rho for a quadratic, and eta inversely scales.
    rate=[]
    for lam in [0.1, .3, 1., 3., 10.]:
        H=torch.eye(2,dtype=torch.float64)*lam; theta=torch.tensor([1.,-.5],dtype=torch.float64)
        def loss(x): return .5*x@H@x
        out=controller_step(theta,loss,eta_prev=.4,eta_max=20.,eta_min=1e-8)
        eta=out[2] * .9 # next_eta is eta/.9; recover accepted eta
        # exact expected accepted eta is rho/lam
        rate.append({'lambda':lam,'predicted_eta':.9/lam,'observed_eta':eta,
                     'eta_times_lambda':eta*lam,'accepted':bool(out[3])})
    # Prediction 3: sufficient-decrease acceptance boundary eta*lambda <= 2(1-c).
    boundary=[]; c=.1
    for lam in [1., 3.]:
        H=torch.eye(1,dtype=torch.float64)*lam; theta=torch.tensor([1.],dtype=torch.float64)
        def loss(x): return .5*x@H@x
        for eta in [1.7/lam, 1.9/lam, 2.1/lam]:
            g=lam*theta; f0=loss(theta).item(); fc=loss(theta-eta*g).item()
            accepted=fc <= f0-c*eta*(g@g).item()+1e-12
            boundary.append({'lambda':lam,'eta_lambda':eta*lam,'predicted_accept':eta*lam <= 2*(1-c), 'observed_accept':accepted})
    return {'curvature_identity':rows,'inverse_scaling':rate,'acceptance_boundary':boundary}

def make_data(n=512):
    rng=np.random.RandomState(SEED)
    x=rng.randn(n,2).astype('float32')
    y=((x[:,0]*x[:,1] + .25*x[:,0] - .15*x[:,1])>0).astype('int64')
    return torch.tensor(x,device=device), torch.tensor(y,device=device)

def net_loss(theta,x,y,d=16):
    a=2*d; b=d; c=2*d
    W1=theta[:a].reshape(d,2); b1=theta[a:a+b]
    W2=theta[a+b:a+b+c].reshape(2,d); b2=theta[a+b+c:a+b+c+2]
    h=torch.tanh(x@W1.T+b1)
    return torch.nn.functional.cross_entropy(h@W2.T+b2,y)

def train_compare(steps=100):
    x,y=make_data(); d=16; n=5*d+2
    torch.manual_seed(SEED)
    init=torch.randn(n,device=device)*.15
    results={}
    for name in ['cosine_sgd','directional']:
        theta=init.clone(); eta=.15; losses=[]; accepted=[]; lrs=[]; t0=time.perf_counter()
        for k in range(steps):
            fn=lambda z: net_loss(z,x,y,d)
            if name=='directional':
                theta,f,eta_next,ok,G,Lh=controller_step(theta,fn,eta,alpha=1.,rho=.9,c=.1,
                                                            eta_min=1e-4,eta_max=2.,max_bt=10)
                eta=eta_next; losses.append(f); accepted.append(ok); lrs.append(eta)
            else:
                th=theta.detach().requires_grad_(True); f=fn(th); g,=torch.autograd.grad(f,th)
                lr=.15*.5*(1+math.cos(math.pi*k/steps)); theta=(th-lr*g).detach()
                losses.append(fn(theta).item()); accepted.append(1); lrs.append(lr)
        with torch.no_grad():
            pred=net_loss(theta,x,y,d).item()
            W1=theta[:2*d].reshape(d,2); b1=theta[2*d:3*d]; W2=theta[3*d:5*d].reshape(2,d); b2=theta[5*d:]
            logits=torch.tanh(x@W1.T+b1)@W2.T+b2
            acc=(logits.argmax(1)==y).float().mean().item()
        results[name]={'final_loss':pred,'accuracy':acc,'loss_at_10':losses[9],
                       'loss_at_50':losses[49],'accepted_fraction':float(np.mean(accepted)),
                       'mean_lr':float(np.mean(lrs)),'wall_seconds':time.perf_counter()-t0}
    return results

if __name__=='__main__':
    out={'device':device,'seed':SEED,'toy':quadratic_sweep(),'training':train_compare()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))