import json, math, os, random, time
from pathlib import Path
import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

SEED = 645
OUT = Path('results.json')

def seed_all(seed=SEED):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)


def harmonic_check():
    # H(q,p)=1/2(q^2+p^2), exact leapfrog versus explicit Euler.
    h = 0.5; n = 2000
    q0, p0 = 1.0, 0.0
    ql, pl, qe, pe = q0, p0, q0, p0
    hl, he = [], []
    for _ in range(n):
        hl.append(0.5*(ql*ql+pl*pl)); he.append(0.5*(qe*qe+pe*pe))
        pl -= 0.5*h*ql
        ql += h*pl
        pl -= 0.5*h*ql
        # explicit Euler for dq/dt=p, dp/dt=-q
        qe, pe = qe + h*pe, pe - h*qe
    hl, he = np.asarray(hl), np.asarray(he)
    return {
        'step': h, 'steps': n,
        'leapfrog_energy_start': float(hl[0]),
        'leapfrog_energy_end': float(hl[-1]),
        'leapfrog_energy_abs_drift': float(abs(hl[-1]-hl[0])),
        'leapfrog_energy_range': float(hl.max()-hl.min()),
        'euler_energy_start': float(he[0]),
        'euler_energy_end': float(he[-1]),
        'euler_energy_growth_factor': float(he[-1]/he[0]),
        'euler_energy_range': float(he.max()-he.min()),
        'claim_observed': bool((hl.max()-hl.min()) < 0.01 and he[-1] > 100*he[0])
    }

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

def make_data():
    x, y = load_digits(return_X_y=True)
    x = StandardScaler().fit_transform(x).astype('float32')
    xtr, xva, ytr, yva = train_test_split(x, y, test_size=.25, random_state=SEED, stratify=y)
    return (torch.tensor(xtr), torch.tensor(ytr, dtype=torch.long),
            torch.tensor(xva), torch.tensor(yva, dtype=torch.long))

def evaluate(model, xv, yv):
    model.eval()
    with torch.no_grad():
        z = model(xv); loss = nn.functional.cross_entropy(z, yv).item()
        acc = (z.argmax(1) == yv).float().mean().item()
    model.train(); return loss, acc

def batches(x, y, bs=128, seed=0):
    g = torch.Generator().manual_seed(seed)
    ix = torch.randperm(len(x), generator=g)
    for s in range(0, len(x), bs):
        j = ix[s:s+bs]; yield x[j], y[j]

def flat_grads(model):
    return torch.cat([p.grad.detach().reshape(-1) for p in model.parameters()])

def flat_params(model):
    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])

def train(kind, xtr, ytr, xva, yva, grad_budget=240):
    # CPU is deliberately used for robust shared-machine reproducibility.
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = MLP().to(device); xtr=xtr.to(device); ytr=ytr.to(device); xva=xva.to(device); yva=yva.to(device)
    if kind == 'sgd': opt = torch.optim.SGD(model.parameters(), lr=0.08, momentum=0.9)
    elif kind == 'adamw': opt = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.0)
    params = list(model.parameters())
    velocity = [torch.zeros_like(p) for p in params]
    mass = 1.0; h = 0.08
    losses=[]; grad_norms=[]; evals=0; step=0
    t0=time.time()
    while evals < grad_budget:
        # Fixed deterministic cycling through reshuffled minibatches.
        for xb, yb in batches(xtr, ytr, 128, SEED+step//100):
            if evals >= grad_budget: break
            if kind == 'leapfrog':
                model.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb), yb); loss.backward()
                g1=flat_grads(model); grad_norms.append(float(g1.norm().cpu()))
                with torch.no_grad():
                    for p, v in zip(params, velocity): p.add_(h*v/mass); v.sub_(0.5*h*p.grad)
                # Correct kick-drift ordering requires first kick before drift; undo/restate compactly:
                # position was advanced using old v, so apply kick ordering from saved old state below is avoided
                # by using the equivalent implementation in the next branch's explicit state correction.
                # This branch is replaced immediately by a clean two-kick update using old parameters.
                raise RuntimeError('internal leapfrog branch should not execute')
            else:
                opt.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb), yb); loss.backward()
                grad_norms.append(float(flat_grads(model).norm().detach().cpu())); opt.step(); evals += 1; step += 1
            if step % 40 == 0: losses.append((step, evals, float(loss.detach().cpu())))
    return {}

def train_leapfrog(xtr,ytr,xva,yva,grad_budget=240):
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model=MLP().to(device); xtr=xtr.to(device); ytr=ytr.to(device); xva=xva.to(device); yva=yva.to(device)
    params=list(model.parameters()); vel=[torch.zeros_like(p) for p in params]; h=.08; evals=0; step=0; gn=[]; trace=[]; t0=time.time()
    while evals+1 <= grad_budget:
        for xb,yb in batches(xtr,ytr,128,SEED+step//100):
            if evals+1>grad_budget: break
            model.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb),yb); loss.backward(); gn.append(float(flat_grads(model).norm().cpu())); evals+=1
            with torch.no_grad():
                for p,v in zip(params,vel): v.sub_(.5*h*p.grad); p.add_(h*v)
            model.zero_grad(set_to_none=True); loss2=nn.functional.cross_entropy(model(xb),yb); loss2.backward(); gn.append(float(flat_grads(model).norm().cpu())); evals+=1
            with torch.no_grad():
                for p,v in zip(params,vel): v.sub_(.5*h*p.grad)
            step+=1
            if step%20==0: trace.append((step,evals,float(loss2.detach().cpu())))
            if evals>=grad_budget: break
    vl,va=evaluate(model,xva,yva)
    return {'steps':step,'gradient_evals':evals,'final_train_loss':float(loss2.detach().cpu()),'val_loss':vl,'val_accuracy':va,'mean_grad_norm':float(np.mean(gn)),'time_sec':time.time()-t0,'trace':trace}

def train_standard(kind,xtr,ytr,xva,yva,grad_budget=240):
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'); model=MLP().to(device)
    xtr=xtr.to(device);ytr=ytr.to(device);xva=xva.to(device);yva=yva.to(device); params=list(model.parameters())
    opt=torch.optim.SGD(params,lr=.08,momentum=.9) if kind=='sgd' else torch.optim.AdamW(params,lr=.003,weight_decay=0.)
    evals=step=0;gn=[];trace=[];t0=time.time()
    while evals<grad_budget:
        for xb,yb in batches(xtr,ytr,128,SEED+step//100):
            if evals>=grad_budget: break
            opt.zero_grad(set_to_none=True); loss=nn.functional.cross_entropy(model(xb),yb); loss.backward(); gn.append(float(flat_grads(model).norm().cpu())); opt.step(); evals+=1;step+=1
            if step%40==0:trace.append((step,evals,float(loss.detach().cpu())))
    vl,va=evaluate(model,xva,yva)
    return {'steps':step,'gradient_evals':evals,'final_train_loss':float(loss.detach().cpu()),'val_loss':vl,'val_accuracy':va,'mean_grad_norm':float(np.mean(gn)),'time_sec':time.time()-t0,'trace':trace}

def main():
    seed_all(); toy=harmonic_check(); xtr,ytr,xva,yva=make_data()
    results={'harmonic_check':toy,'device':str(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))}
    for k in ['sgd','adamw','leapfrog']:
        seed_all(); results[k]=train_leapfrog(xtr,ytr,xva,yva) if k=='leapfrog' else train_standard(k,xtr,ytr,xva,yva)
    OUT.write_text(json.dumps(results,indent=2)); print(json.dumps(results,indent=2))
if __name__=='__main__': main()