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

SEED = int(os.environ.get('SEED', '3126'))
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type == 'cuda': torch.zeros(1, device=device)
except Exception:
    device = torch.device('cpu')


def polynomial_check():
    out = {}
    E0, c, dt, n = 2.0, 0.25, 1e-3, 10000
    for m in (0, 1, 2):
        E = E0; hist = [E]; chats = []
        for _ in range(n):
            old = E
            E = max(0.0, E - dt*c*E**(1+m))
            chats.append((old-E)/(dt*old**(1+m)))
            hist.append(E)
        t = np.arange(n+1)*dt
        exact = E0*np.exp(-c*t) if m == 0 else (E0**(-m)+m*c*t)**(-1/m)
        ix = np.arange(n//2+1, n+1)
        slope = np.polyfit(np.log(t[ix]), np.log(np.maximum(np.asarray(hist)[ix],1e-30)), 1)[0]
        out[str(m)] = {
            'relative_final_error': float(abs(E-exact[-1])/exact[-1]),
            'chat_mean': float(np.mean(chats)), 'chat_std': float(np.std(chats)),
            'loglog_slope': float(slope), 'expected_slope': None if m == 0 else -1/m}
    return out


class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.seq = torch.nn.Sequential(torch.nn.Linear(64,32), torch.nn.Tanh(), torch.nn.Linear(32,10))
    def forward(self, x): return self.seq(x)


def make_data():
    x, y = load_digits(return_X_y=True)
    x = ((x / 16.0) - 0.5).astype('float32')
    xa, xb, ya, yb = train_test_split(x, y, test_size=0.25, random_state=SEED, stratify=y)
    return torch.tensor(xa), torch.tensor(ya), torch.tensor(xb), torch.tensor(yb)


def train(kind, m=1, steps=500, batch=64):
    torch.manual_seed(SEED)
    x, y, xv, yv = make_data()
    model = Net().to(device)
    lossfn = torch.nn.CrossEntropyLoss()
    p = list(model.parameters())
    eta = 0.12 if kind == 'baseline' else 0.12
    eta_min, eta_max = 1e-4, 0.5
    cstar, alpha, smooth = {0: 0.003, 1: 0.001, 2: 0.0005}[m], 0.20, 0.85
    floor = 1e-5
    rng = np.random.default_rng(SEED)
    # Fixed monitoring batch makes E comparisons meaningful and low-noise.
    mi = torch.arange(min(128, len(xv)))
    xm, ym = xv[mi].to(device), yv[mi].to(device)
    rejects = 0; accepted = 0; cema = None; losses = []
    for step in range(steps):
        ii = torch.tensor(rng.integers(0, len(x), size=batch))
        xb, yb = x[ii].to(device), y[ii].to(device)
        model.zero_grad(set_to_none=True)
        loss = lossfn(model(xb), yb); loss.backward()
        old_state = {k:v.detach().clone() for k,v in model.state_dict().items()}
        with torch.no_grad():
            oldE = float(lossfn(model(xm), ym).item())
            for q in p: q.add_(q.grad, alpha=-eta)
            trialE = float(lossfn(model(xm), ym).item())
        if kind == 'baseline':
            accepted += 1
        elif trialE > oldE:
            rejects += 1
            model.load_state_dict(old_state)
            eta = max(eta_min, eta * 0.5)
        else:
            accepted += 1
            E = max(oldE - floor, floor)
            chat = (oldE - trialE) / max(E**(1+m), 1e-12)
            cema = chat if cema is None else smooth*cema + (1-smooth)*chat
            ratio = max(cema, 1e-12)/cstar
            eta = float(np.clip(eta * ratio**alpha, eta_min, eta_max))
        with torch.no_grad(): losses.append(float(lossfn(model(xm), ym).item()))
    with torch.no_grad(): test_loss=float(lossfn(model.to(device)(xv.to(device)), yv.to(device)).item())
    return {'final_monitor_loss': losses[-1], 'test_loss': test_loss, 'rejects': rejects,
            'acceptance_rate': accepted/steps, 'final_eta': eta,
            'loss_first': losses[0], 'loss_min': min(losses)}


def main():
    result = {'device': str(device), 'math_check': polynomial_check(), 'training': {}}
    for name, m in [('baseline', 0), ('m0', 0), ('m1', 1), ('m2', 2)]:
        result['training'][name] = train('baseline' if name == 'baseline' else 'controller', m=m)
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()