import json
import math
import random
from pathlib import Path

import numpy as np
import torch
from torch import nn

SEED = 587
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)


class GRUSequence(nn.Module):
    def __init__(self, hidden=24):
        super().__init__()
        self.gru = nn.GRU(1, hidden, batch_first=True)
        self.readout = nn.Linear(hidden, 1)

    def forward(self, u):
        # h_all includes h_1,...,h_T; h0 is explicitly zero.
        out, _ = self.gru(u)
        h0 = torch.zeros(u.shape[0], 1, out.shape[-1], device=u.device)
        h_all = torch.cat([h0, out], dim=1)
        y = self.readout(out).squeeze(-1)
        return y, h_all


def passivity_loss(u, y, h_all, gamma=0.10, incremental=None, eta=0.0):
    # u: B,T,1; y: B,T; h_all: B,T+1,H
    r = (h_all[:, 1:].pow(2).sum(-1) + y.pow(2)
         - h_all[:, :-1].pow(2).sum(-1) - u.squeeze(-1).pow(2))
    lp = torch.relu(r).mean() + gamma * torch.relu(r.sum(dim=1)).mean()
    inc_lp = torch.zeros((), device=u.device)
    if incremental is not None and eta:
        u2, y2, h2 = incremental
        du = u2 - u
        dy = y2 - y
        dh = h2 - h_all
        ri = (dh[:, 1:].pow(2).sum(-1) + dy.pow(2)
              - dh[:, :-1].pow(2).sum(-1) - du.squeeze(-1).pow(2))
        inc_lp = torch.relu(ri).mean()
        lp = lp + eta * inc_lp
    stats = {
        "max_positive_residual": float(torch.relu(r).max().detach().cpu()),
        "mean_cumulative_residual": float(r.sum(1).mean().detach().cpu()),
        "mean_hidden_norm": float(h_all[:, 1:].norm(dim=-1).mean().detach().cpu()),
        "penalty": float(lp.detach().cpu()),
        "incremental_penalty": float(inc_lp.detach().cpu()),
    }
    return lp, stats


def math_check():
    # Directly check sum_t r_t = ||h_T||^2 + sum ||y||^2 - ||h_0||^2 - sum ||u||^2.
    g = torch.Generator().manual_seed(SEED)
    B, T, H = 7, 13, 5
    u = torch.randn(B, T, 1, generator=g)
    h = torch.randn(B, T + 1, H, generator=g)
    y = torch.randn(B, T, generator=g)
    r = h[:, 1:].pow(2).sum(-1) + y.pow(2) - h[:, :-1].pow(2).sum(-1) - u.squeeze(-1).pow(2)
    lhs = r.sum(1)
    rhs = h[:, -1].pow(2).sum(-1) + y.pow(2).sum(1) - h[:, 0].pow(2).sum(-1) - u.squeeze(-1).pow(2).sum(1)
    err = float((lhs - rhs).abs().max())
    # A deliberately amplifying scalar trajectory has positive residuals.
    hs = [torch.tensor([[1.0]])]
    us, ys = [], []
    for _ in range(8):
        us.append(torch.zeros(1, 1))
        ys.append(torch.zeros(1))
        hs.append(1.15 * hs[-1])
    rr = []
    for t in range(8):
        rr.append(hs[t+1].pow(2).sum() - hs[t].pow(2).sum())
    positive_fraction = float((torch.stack(rr) > 0).float().mean())
    return {"telescoping_max_abs_error": err, "amplifying_positive_residual_fraction": positive_fraction}


def batch_data(n, T, seed):
    # Long-context signal: classify whether the sum of random +/-1 inputs is positive.
    gen = torch.Generator().manual_seed(seed)
    x = torch.randint(0, 2, (n, T, 1), generator=gen).float() * 2 - 1
    # Add a weak final tie breaker so labels are balanced and deterministic.
    sums = x.sum(dim=1).squeeze(-1)
    labels = (sums > 0).float()
    ties = sums == 0
    labels[ties] = (x[ties, -1, 0] > 0).float()
    return x, labels


def run_model(reg_lambda, paired, seed=SEED):
    seed_all(seed)
    model = GRUSequence().to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    train_u, train_y = batch_data(96, 32, 1001)
    train_u, train_y = train_u.to(DEVICE), train_y.to(DEVICE)
    last = {}
    for step in range(260):
        # Fresh small batches are unnecessary here; fixed data makes the comparison cheap/reproducible.
        if paired:
            noise = 0.05 * torch.randn_like(train_u)
            u2 = train_u + noise
            y2, h2 = model(u2)
        y, h = model(train_u)
        task = nn.functional.binary_cross_entropy_with_logits(y[:, -1], train_y)
        inc = (u2, y2, h2) if paired else None
        pl, st = passivity_loss(train_u, y, h, gamma=0.10, incremental=inc, eta=0.10 if paired else 0.0)
        loss = task + reg_lambda * pl
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        opt.step()
        last = {"task": float(task.detach().cpu()), **st}

    def evaluate(T, seed_offset):
        u, target = batch_data(512, T, 4000 + seed_offset)
        u, target = u.to(DEVICE), target.to(DEVICE)
        with torch.no_grad():
            y, h = model(u)
            acc = ((torch.sigmoid(y[:, -1]) > .5) == target.bool()).float().mean()
            _, st = passivity_loss(u, y, h, gamma=0.10)
            # A finite input perturbation amplification at the final state.
            u2 = u + 0.02 * torch.randn_like(u)
            _, h2 = model(u2)
            amp = ((h2[:, -1] - h[:, -1]).norm(dim=-1) /
                   (u2 - u).norm(dim=(1, 2)).clamp_min(1e-8)).mean()
        return {"accuracy": float(acc.cpu()), "amplification": float(amp.cpu()), **st}

    return {"train": last, "test_T32": evaluate(32, 32), "test_T160": evaluate(160, 160)}


def main():
    seed_all()
    math_result = math_check()
    results = {
        "device": DEVICE,
        "math_check": math_result,
        "baseline": run_model(0.0, False, SEED),
        "passivity_lambda_0.01": run_model(0.01, False, SEED),
        "passivity_lambda_0.01_incremental": run_model(0.01, True, SEED),
    }
    Path("results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        # CUDA can fail in a shared environment; retry entirely on CPU.
        if DEVICE == "cuda":
            print("CUDA failed, retrying on CPU:", repr(e))
            torch.cuda.empty_cache()
            DEVICE = "cpu"
            main()
        else:
            raise