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

SEED = 7
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)
try:
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
except Exception:
    DEVICE = torch.device("cpu")

class DissipativeDrift(nn.Module):
    """b=-mu*a+Lf*tanh(a+h(t,x)); action Jacobian norm <= mu+Lf in general,
    and the residual has action Lipschitz constant Lf, as required by the claim."""
    def __init__(self, state_dim=2, action_dim=2, mu=0.8, lf=0.35):
        super().__init__()
        self.mu, self.lf = mu, lf
        self.center = nn.Sequential(nn.Linear(state_dim + 1, 32), nn.Tanh(),
                                     nn.Linear(32, action_dim))
    def forward(self, t, x, a):
        h = self.center(torch.cat([t, x], dim=-1))
        return -self.mu * a + self.lf * torch.tanh(a + h)

class GenericDrift(nn.Module):
    def __init__(self, state_dim=2, action_dim=2):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(state_dim + 1 + action_dim, 64), nn.Tanh(),
                                 nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, action_dim))
    def forward(self, t, x, a):
        return self.net(torch.cat([t, x, a], dim=-1))

def target_drift(t, x, a):
    c = torch.cat([t, x], -1)
    h = torch.cat([0.55 * torch.sin(c[:, :1] + c[:, 1:2]),
                   0.45 * torch.cos(c[:, :1] - c[:, 2:3])], -1)
    return -0.8 * a + h + 0.12 * torch.tanh(2.0 * a)

def sample(n, device):
    t = torch.rand(n, 1, device=device)
    x = torch.randn(n, 2, device=device)
    a = 1.5 * torch.randn(n, 2, device=device)
    return t, x, a

def train(model, steps=1200, batch=128):
    model.to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    for _ in range(steps):
        t, x, a = sample(batch, DEVICE)
        loss = ((model(t, x, a) - target_drift(t, x, a)) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        t, x, a = sample(4096, DEVICE)
        mse = ((model(t, x, a) - target_drift(t, x, a)) ** 2).mean().item()
    return mse

def verify(model, mu=0.8, lf=0.35, n=10000, residual=False):
    # Same t,x for each pair, exactly matching the action-coordinate claim.
    t, x, a = sample(n, DEVICE)
    ap = a + 0.5 * torch.randn_like(a)
    with torch.no_grad():
        db = model(t, x, a) - model(t, x, ap)
        da = a - ap
        lhs = (db * da).sum(1)
        rhs = -(mu - lf) * (da * da).sum(1)
        ratio = (lhs / (da * da).sum(1)).cpu().numpy()
        violations = (lhs - rhs > 1e-6).float().mean().item()
        # Residual finite-difference Lipschitz estimate, only for the proposed branch.
        if residual:
            h = model.center(torch.cat([t, x], -1))
            fp = lf * torch.tanh(ap + h)
            f = lf * torch.tanh(a + h)
            lip = (fp - f).norm(dim=1) / (ap - a).norm(dim=1).clamp_min(1e-8)
            max_lip = float(lip.max())
        else:
            max_lip = None
    return {"worst_one_sided_ratio": float(ratio.max()),
            "theoretical_bound": -(mu-lf),
            "violation_fraction": violations,
            "max_residual_fd_lipschitz": max_lip}

def rollout(model, steps, n=512, dt=0.08, sigma=0.18):
    model.eval()
    t0 = torch.zeros(n, 1, device=DEVICE)
    x = torch.randn(n, 2, device=DEVICE)
    a = 2.5 * torch.randn(n, 2, device=DEVICE)
    with torch.no_grad():
        for k in range(steps):
            t = torch.full((n, 1), k * dt, device=DEVICE)
            a = a + dt * model(t, x, a) + sigma * math.sqrt(dt) * torch.randn_like(a)
    return float(a.norm(dim=1).mean().cpu()), float(a.norm(dim=1).std().cpu())

def main():
    # Training is deliberately identical; the comparison is a toy drift-fitting proxy.
    generic = GenericDrift()
    dissip = DissipativeDrift()
    generic_mse = train(generic)
    dissip_mse = train(dissip)
    check = verify(dissip, residual=True)
    generic_check = verify(generic)
    results = {"device": str(DEVICE), "generic_validation_mse": generic_mse,
               "dissipative_validation_mse": dissip_mse, "math_check_dissipative": check, "math_check_generic_control": generic_check, "rollouts": {}}
    for steps in (8, 20, 50):
        results["rollouts"][str(steps)] = {"generic": rollout(generic, steps),
                                            "dissipative": rollout(dissip, steps)}
    Path("results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    main()