import json, math, time
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED = 281
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def hankel(values, k=3):
    # values has length 2k-1 and gradients are retained.
    return torch.stack([torch.stack([values[i + j] for j in range(k)]) for i in range(k)])


def hankel_penalty(values, k=3, eps=1e-4):
    H = hankel(values, k)
    eig = torch.linalg.eigvalsh(H)
    # Smooth enough almost everywhere and directly measures PSD violations.
    violation = torch.relu(eps - eig)
    return (violation ** 2).mean(), eig.detach().cpu().numpy()


def pairwise_logconvex_penalty(values, eps=1e-6):
    # For a positive log-convex sequence: f_r^2 <= f_{r-1} f_{r+1}.
    curv = torch.log(values[2:] + eps) - 2 * torch.log(values[1:-1] + eps) + torch.log(values[:-2] + eps)
    return torch.relu(-curv).pow(2).mean()


def analytic_check():
    # f(a)=sum_s w_s exp(-s a), so H is a Gram matrix with vectors exp(-s*i*Delta).
    a, delta = 0.7, 0.35
    s = torch.tensor([0.15, 0.8, 1.7], dtype=torch.float64)
    w = torch.tensor([0.7, 0.2, 0.1], dtype=torch.float64)
    vals = torch.stack([(w * torch.exp(-s * (a + r * delta))).sum() for r in range(5)])
    H = torch.stack([torch.stack([vals[i+j] for j in range(3)]) for i in range(3)])
    eig = torch.linalg.eigvalsh(H).numpy()
    # Deliberately perturb one middle sample to demonstrate the violation detector.
    bad = vals.clone(); bad[2] *= 1.8
    badH = torch.stack([torch.stack([bad[i+j] for j in range(3)]) for i in range(3)])
    bad_eig = torch.linalg.eigvalsh(badH).numpy()
    return {"valid_values": vals.numpy().tolist(), "valid_eigenvalues": eig.tolist(),
            "valid_min_eigenvalue": float(eig.min()), "perturbed_eigenvalues": bad_eig.tolist(),
            "perturbed_min_eigenvalue": float(bad_eig.min()),
            "claim_observed": bool(eig.min() >= -1e-10 and bad_eig.min() < -1e-6)}


class ResponseNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(1, 24), nn.Tanh(), nn.Linear(24, 24), nn.Tanh(), nn.Linear(24, 1))
    def forward(self, a):
        a = a.reshape(-1, 1)
        return torch.nn.functional.softplus(self.net(a).squeeze(-1)) + 1e-4


def train(kind, train_a, train_y, test_a, test_y, steps=700):
    torch.manual_seed(SEED + {"none":0, "pairwise":1, "hankel":2}[kind])
    model = ResponseNet().to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=0.012)
    # Five nearby evaluations are needed for K=3.
    delta = 0.075
    t0 = time.perf_counter()
    last_eigs = None
    for step in range(steps):
        opt.zero_grad(set_to_none=True)
        pred = model(train_a)
        loss = ((pred - train_y) ** 2).mean()
        if kind != "none":
            anchors = train_a[:, None]
            nearby = model((anchors[:, 0, None] + delta * torch.arange(5, device=DEVICE)[None, :]).reshape(-1)).reshape(-1, 5)
            # Average penalty over each anchor; this evaluates the same response head nearby.
            if kind == "hankel":
                H = torch.stack([nearby[:, i:i+3] for i in range(3)], dim=1)
                eig_batch = torch.linalg.eigvalsh(H)
                reg = torch.relu(1e-4 - eig_batch).pow(2).mean()
                last_eigs = eig_batch.detach().cpu().numpy()
            else:
                curv = torch.log(nearby[:, 2:] + 1e-6) - 2 * torch.log(nearby[:, 1:-1] + 1e-6) + torch.log(nearby[:, :-2] + 1e-6)
                reg = torch.relu(-curv).pow(2).mean()
            loss = loss + 3.0 * reg
        loss.backward(); opt.step()
    elapsed = time.perf_counter() - t0
    with torch.no_grad():
        pred_test = model(test_a)
        # Dense-grid response for oscillation and Hankel violations.
        grid = torch.linspace(0.15, 3.0, 121, device=DEVICE)
        gv = model(grid)
        second = gv[2:] - 2*gv[1:-1] + gv[:-2]
        oscillation = float(torch.abs(second).mean().cpu())
        dense_eigs = []
        for i in range(len(grid)-4):
            _, e = hankel_penalty(gv[i:i+5], 3, 1e-4); dense_eigs.append(e)
        dense_eigs = np.asarray(dense_eigs)
    return {"test_mse": float(((pred_test-test_y)**2).mean().cpu()),
            "train_mse": float(((model(train_a)-train_y)**2).mean().cpu()),
            "mean_abs_second_difference": oscillation,
            "min_dense_hankel_eigenvalue": float(dense_eigs.min()),
            "fraction_dense_psd_violations": float((dense_eigs < -1e-7).any(axis=1).mean()),
            "seconds": elapsed, "relative_overhead": None}


def main():
    check = analytic_check()
    # Sparse, mildly noisy observations make the extrapolating shape meaningful.
    true_s, true_w = np.array([0.2, 0.9, 2.0]), np.array([0.55, 0.3, 0.15])
    def f(x): return (true_w[None,:] * np.exp(-x[:,None]*true_s[None,:])).sum(1)
    train_x = np.linspace(0.25, 2.35, 18).astype("float32")
    test_x = np.linspace(0.15, 3.0, 121).astype("float32")
    rng = np.random.RandomState(SEED)
    y = (f(train_x) + rng.normal(0, 0.012, len(train_x))).clip(1e-3).astype("float32")
    ty = f(test_x).astype("float32")
    ta, va = torch.tensor(train_x, device=DEVICE), torch.tensor(y, device=DEVICE)
    te, vy = torch.tensor(test_x, device=DEVICE), torch.tensor(ty, device=DEVICE)
    results = {}
    for kind in ["none", "pairwise", "hankel"]:
        results[kind] = train(kind, ta, va, te, vy)
    base = results["none"]["seconds"]
    for r in results.values(): r["relative_overhead"] = r["seconds"] / base
    out = {"device": DEVICE, "analytic_check": check, "results": results,
           "setup": {"steps": 700, "train_points": 18, "test_points": 121, "noise_std": 0.012, "K": 3}}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
        if DEVICE == "cuda":
            print("CUDA failed; rerun with CPU", repr(e))
            torch.cuda.empty_cache()
            DEVICE = "cpu"
            main()
        else: raise