import json, 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 = 591
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.set_per_process_memory_fraction(0.16)
except Exception:
    device = torch.device('cpu')


def math_check():
    # x'=-lambda*x: compare explicit Euler with denominator Euler.
    lam, alpha = 1.0, 0.5
    hs = np.array([0.5, 1., 2., 4., 8., 32.])
    q = hs / (1 + alpha * hs)
    euler_amp = np.abs(1 - hs * lam)
    ns_amp = np.abs(1 - q * lam)
    return {
        'h': hs.tolist(), 'q': q.tolist(),
        'euler_amplification': euler_amp.tolist(),
        'nonstandard_amplification': ns_amp.tolist(),
        'q_bound_1_over_alpha': 1 / alpha,
        'nonstandard_nonexpansive_all_tested': bool(np.all(ns_amp <= 1 + 1e-12)),
        'euler_nonexpansive_all_tested': bool(np.all(euler_amp <= 1 + 1e-12))
    }


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


class ResidualModel(nn.Module):
    def __init__(self, d, blocks, h, kind, alpha=0.5, fixed_iters=6):
        super().__init__()
        self.blocks = nn.ModuleList([MLP(d) for _ in range(blocks)])
        self.h, self.kind, self.alpha, self.fixed_iters = h, kind, alpha, fixed_iters
        self.head = nn.Linear(d, 10)
        self.last_stats = {'max_activation': 0., 'max_residual': 0., 'failed_solves': 0}

    def forward(self, x):
        max_act, max_res, failures = 0., 0., 0
        for f in self.blocks:
            if self.kind == 'baseline':
                y = x + self.h * f(x)
                residual = (y - x).norm(dim=1).mean().item()
            else:
                # Two-stage explicit Heun tableau, with denominator q replacing h.
                q = self.h / (1 + self.alpha * self.h)
                g1 = x
                g2 = x + q * f(g1)
                for _ in range(self.fixed_iters - 1):
                    old = g2
                    g2 = x + q * f(g1)  # stage-2 is explicit in this concrete SSP tableau
                    err = (g2 - old).norm(dim=1).mean().item()
                    max_res = max(max_res, err)
                y = x + q * 0.5 * (f(g1) + f(g2))
                residual = (y - x).norm(dim=1).mean().item()
                if not torch.isfinite(y).all(): failures += 1
            x = y
            max_act = max(max_act, x.detach().norm(dim=1).max().item())
            max_res = max(max_res, residual)
        self.last_stats = {'max_activation': max_act, 'max_residual': max_res,
                           'failed_solves': failures}
        return self.head(x)


def run_one(kind, h, Xtr, ytr, Xte, yte, epochs=18):
    torch.manual_seed(SEED + int(h * 10) + (0 if kind == 'baseline' else 100))
    model = ResidualModel(Xtr.shape[1], blocks=16, h=h, kind=kind).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=2e-3)
    loss_fn = nn.CrossEntropyLoss()
    max_grad, nan = 0., False
    t0 = time.time()
    for _ in range(epochs):
        perm = torch.randperm(len(Xtr), device=device)
        for ix in perm.split(128):
            opt.zero_grad(set_to_none=True)
            out = model(Xtr[ix]); loss = loss_fn(out, ytr[ix])
            if not torch.isfinite(loss): nan = True; break
            loss.backward()
            grad = torch.nn.utils.clip_grad_norm_(model.parameters(), 10.)
            max_grad = max(max_grad, float(grad))
            opt.step()
        if nan: break
    with torch.no_grad():
        logits = model(Xte); test_loss = float(loss_fn(logits, yte)); acc = float((logits.argmax(1) == yte).float().mean())
    return {'loss': test_loss, 'accuracy': acc, 'max_activation': model.last_stats['max_activation'],
            'max_update_proxy': model.last_stats['max_residual'], 'max_grad': max_grad,
            'failed_solves': model.last_stats['failed_solves'], 'nan': nan,
            'seconds': time.time() - t0, 'parameters': sum(p.numel() for p in model.parameters())}


def main():
    data = load_digits()
    X = StandardScaler().fit_transform(data.data).astype('float32')
    y = data.target.astype('int64')
    Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=.25, random_state=SEED, stratify=y)
    tensors = [torch.tensor(z, device=device) for z in (Xtr, ytr, Xte, yte)]
    Xtr, ytr, Xte, yte = tensors
    results = {'device': str(device), 'math_check': math_check(), 'runs': {}}
    for h in (0.5, 1., 2., 4.):
        for kind in ('baseline', 'idea'):
            results['runs'][f'{kind}_h{h}'] = run_one(kind, h, Xtr, ytr, Xte, yte)
            print(kind, h, results['runs'][f'{kind}_h{h}'])
    Path('results.json').write_text(json.dumps(results, indent=2))
    print('Wrote results.json')

if __name__ == '__main__': main()