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

SEED = 624
M = 4
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def seed_all(seed=SEED):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def variance_check():
    rng = np.random.default_rng(SEED)
    sigma = 1.7
    n = 300000
    rows = []
    for m in [1, 2, 4, 8, 16]:
        errors = rng.normal(0.0, sigma, size=(n, m))
        observed = float(np.var(errors.mean(axis=1)))
        predicted = sigma * sigma / m
        rows.append({"M": m, "observed": observed, "predicted": predicted,
                     "ratio": observed / predicted})
    return rows


def oscillator_step(x, dt=0.08):
    q, v = x[:, 0], x[:, 1]
    q_next = q + dt * v
    v_next = v + dt * (-0.8 * q - 0.15 * v + 0.10 * q ** 3)
    return torch.stack([q_next, v_next], dim=1)


def make_data(n, seed):
    generator = torch.Generator().manual_seed(seed)
    x = torch.empty(n, 2).uniform_(-1.2, 1.2, generator=generator)
    y = oscillator_step(x)
    return x, y


class PhaseNet(nn.Module):
    def __init__(self, m=M):
        super().__init__()
        self.feature = nn.Sequential(
            nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh()
        )
        self.heads = nn.ModuleList([nn.Linear(32, 2) for _ in range(m)])

    def forward(self, x):
        h = self.feature(x)
        return torch.stack([head(h) for head in self.heads], dim=1)


def train(mode, x_train, y_train, phase_offsets, steps=500):
    model = PhaseNet(M).to(DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=2e-3)
    mse = nn.MSELoss()
    x_train, y_train = x_train.to(DEVICE), y_train.to(DEVICE)
    batch_size = 128
    for step in range(steps):
        idx = torch.randint(0, len(x_train), (batch_size,), device=DEVICE)
        pred = model(x_train[idx])
        target = y_train[idx]
        # Simulate cyclic phase alignment by assigning each head a small,
        # fixed offset in its phase-specific target.
        targets = target[:, None, :].expand(-1, M, -1) + phase_offsets[None, :, :]
        losses = ((pred - targets) ** 2).mean(dim=(0, 2))
        centroid_loss = ((pred.mean(dim=1) - target) ** 2).mean()
        if mode == "single":
            loss = losses[0]
        elif mode == "centroid":
            loss = centroid_loss
        elif mode == "robust":
            worst = losses.max()
            # Output-space convex mixture, choosing the worst vertex exactly.
            loss = centroid_loss + 0.5 * worst
        else:
            raise ValueError(mode)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
    return model


def evaluate(model, x, y, phase_offsets):
    model.eval()
    with torch.no_grad():
        pred = model(x.to(DEVICE))
        target = y.to(DEVICE)
        targets = target[:, None, :].expand(-1, M, -1) + phase_offsets[None, :, :]
        vertex_mse = ((pred - targets) ** 2).mean(dim=(0, 2))
        centroid_mse = ((pred.mean(dim=1) - target) ** 2).mean()
        spread = pred.var(dim=1, unbiased=False).mean()
    return {
        "centroid_rmse": float(torch.sqrt(centroid_mse).cpu()),
        "max_vertex_rmse": float(torch.sqrt(vertex_mse.max()).cpu()),
        "mean_vertex_rmse": float(torch.sqrt(vertex_mse.mean()).cpu()),
        "vertex_loss_spread": float(vertex_mse.max().cpu() - vertex_mse.min().cpu()),
        "prediction_spread": float(spread.cpu()),
    }


def main():
    seed_all()
    x_train, y_train = make_data(4096, SEED + 1)
    x_test, y_test = make_data(2048, SEED + 2)
    # Small phase-dependent observation shifts make the heads genuinely distinct.
    offsets = torch.tensor([
        [-0.030, 0.020], [0.020, -0.025], [0.015, 0.030], [-0.025, -0.020]
    ], dtype=torch.float32, device=DEVICE)
    results = {"device": DEVICE, "variance_check": variance_check(), "models": {}}
    for mode in ["single", "centroid", "robust"]:
        seed_all(SEED + {"single": 1, "centroid": 2, "robust": 3}[mode])
        model = train(mode, x_train, y_train, offsets)
        results["models"][mode] = evaluate(model, x_test, y_test, offsets)
    print(json.dumps(results, indent=2, sort_keys=True))
    Path("phase_polytope_results.json").write_text(json.dumps(results, indent=2, sort_keys=True))


if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, torch.cuda.CudaError) as exc:
        if DEVICE == "cuda":
            print("CUDA failed; rerun on CPU:", repr(exc))
            torch.cuda.empty_cache()
            torch.cuda.is_available = lambda: False
            main()
        else:
            raise