import json
import math
import random
import numpy as np
import torch
from torch import nn

SEED = 7
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = torch.device('cpu')
DT = 0.08
DIM = 2


def rk4(fun, x, dt=DT):
    k1 = dt * fun(x)
    k2 = dt * fun(x + 0.5*k1)
    k3 = dt * fun(x + 0.5*k2)
    k4 = dt * fun(x + k3)
    return x + (k1 + 2*k2 + 2*k3 + k4) / 6.0


def true_field(x):
    # Stable nonlinear oscillator with state-dependent frequency and damping.
    q, p = x[..., 0], x[..., 1]
    return torch.stack((p, -0.8*q - 0.15*p - 0.18*q**3), dim=-1)


def make_data(ntraj=96, length=72, noise=0.002):
    g = torch.Generator().manual_seed(SEED)
    x = torch.zeros(ntraj, length+1, DIM)
    x[:, 0] = torch.empty(ntraj, DIM).uniform_(-1.3, 1.3, generator=g)
    with torch.no_grad():
        for t in range(length):
            x[:, t+1] = rk4(true_field, x[:, t])
    # Observational noise is applied to training targets, while evaluation uses clean trajectories.
    noisy = x + noise * torch.randn(x.shape, generator=g)
    return x, noisy


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


def rollout(model, x0, horizon):
    xs = [x0]
    x = x0
    for _ in range(horizon):
        x = rk4(model, x)
        xs.append(x)
    return torch.stack(xs, dim=1)


def l1_penalty(model):
    return sum(p.abs().sum() for p in model.parameters())


def train_one_step(train_noisy, epochs=105, batch=24):
    model = Field().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    n = train_noisy.shape[0]
    model.train()
    for ep in range(epochs):
        order = torch.randperm(n)
        for start in range(0, n, batch):
            ids = order[start:start+batch]
            # Random teacher-forced state and its next noisy observation.
            t = torch.randint(0, train_noisy.shape[1]-1, (len(ids),))
            x0 = train_noisy[ids, t]
            target = train_noisy[ids, t+1]
            pred = rk4(model, x0)
            loss = (pred-target).abs().mean() + 1e-6*l1_penalty(model)
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
    return model


def train_progressive(train_noisy, schedule=(1,2,4,8), epochs_per_phase=27, batch=24):
    model = Field().to(device)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    n = train_noisy.shape[0]
    model.train()
    for H in schedule:
        for ep in range(epochs_per_phase):
            order = torch.randperm(n)
            for start in range(0, n, batch):
                ids = order[start:start+batch]
                max_t = train_noisy.shape[1] - 1 - H
                t = torch.randint(0, max_t+1, (len(ids),))
                target = train_noisy[ids[:, None], t[:, None] + torch.arange(1,H+1)[None, :]]
                x = train_noisy[ids, t]
                pred = rollout(model, x, H)[:, 1:]
                loss = (pred-target).abs().mean() + 1e-6*l1_penalty(model)
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
    return model


def evaluate(model, clean, horizons=(8,16,32,64)):
    model.eval(); out={}
    with torch.no_grad():
        for H in horizons:
            pred = rollout(model, clean[:,0], H)[:,1:]
            target = clean[:,1:H+1]
            err = (pred-target).abs().mean().item()
            maxnorm = pred.norm(dim=-1).max().item()
            divergent = (pred.norm(dim=-1).max(dim=1).values > 4.0).float().mean().item()
            out[str(H)] = {'mae':err, 'max_norm':maxnorm, 'divergence_rate':divergent}
    return out


def rk4_sanity():
    # For x'=A x, RK4 error should decrease by ~16 when dt is halved.
    A = torch.tensor([[-0.2, 1.0],[-1.4,-0.3]], dtype=torch.float64)
    x0 = torch.tensor([[1.1,-0.4]], dtype=torch.float64)
    def f(x): return x @ A.T
    # Same physical time, reference from a very fine RK4 integration.
    with torch.no_grad():
        ref=x0.clone()
        for _ in range(40000): ref=rk4(f, ref, 0.00002)
        errs=[]
        for dt, steps in [(0.16,5),(0.08,10),(0.04,20)]:
            y=x0.clone()
            for _ in range(steps): y=rk4(f,y,dt)
            errs.append(float((y-ref).abs().max()))
    ratios=[errs[i]/errs[i+1] for i in range(2)]
    return {'errors_dt_0.16_0.08_0.04':errs, 'halving_error_ratios':ratios, 'passes_fourth_order_signal': all(r>8 for r in ratios)}


def main():
    sanity = rk4_sanity()
    clean, noisy = make_data()
    # Equal optimizer epochs and same architecture; progressive phases expose the model to all horizons.
    baseline = train_one_step(noisy, epochs=225)
    progressive = train_progressive(noisy, schedule=(1,2,4,8), epochs_per_phase=15)
    result = {
        'seed': SEED, 'dt': DT, 'train_trajectories': len(clean),
        'rk4_sanity': sanity,
        'baseline': evaluate(baseline, clean),
        'progressive': evaluate(progressive, clean),
    }
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()