import json, math, random
from pathlib import Path
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

SEED = 1729

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 toy_check():
    # A zero-conditional-mean increment and a history-dependent biased increment.
    rng = np.random.default_rng(SEED)
    n = 200000
    h = rng.normal(size=n)
    noise = rng.normal(size=n)
    # bounded history test functions, including a nonlinear one
    hs = np.stack([np.ones(n), np.tanh(h), h/(1+np.abs(h))], axis=1)
    unbiased = 0.08 * noise
    biased = 0.08 * noise + 0.08 * np.tanh(h)
    norms = []
    for d in (unbiased, biased):
        moments = (d[:, None] * hs).mean(axis=0)
        norms.append(float(np.sqrt(np.sum(moments**2))))
    return {"unbiased_moment_norm": norms[0], "biased_moment_norm": norms[1],
            "ratio_biased_to_unbiased": norms[1] / max(norms[0], 1e-12),
            "passed": bool(norms[1] > 5 * norms[0])}


def make_data(n, steps=6, obs_dim=8, seed=SEED):
    rng = np.random.default_rng(seed)
    y = rng.integers(0, 3, size=n)
    # Each revealed observation is a noisy class-dependent vector; later steps add evidence.
    means = np.array([[1.0, 0.0, 0.0, 0.0, .3, -.2, 0, 0],
                      [0.0, 1.0, 0.0, 0.0, -.2, .3, 0, 0],
                      [0.0, 0.0, 1.0, 0.0, 0, 0, .3, -.2]], dtype=np.float32)
    x = means[y, None, :] + rng.normal(0, 1.15, size=(n, steps, obs_dim)).astype(np.float32)
    return torch.tensor(x), torch.tensor(y, dtype=torch.long)


class SeqClassifier(nn.Module):
    def __init__(self, obs_dim=8, hidden=32, classes=3):
        super().__init__()
        self.gru = nn.GRU(obs_dim, hidden, batch_first=True)
        self.head = nn.Linear(hidden, classes)
    def forward(self, x):
        z, _ = self.gru(x)
        return self.head(z)


def mart_loss(logits):
    p = logits.softmax(-1)
    d = p[:, 1:] - p[:, :-1]                 # B,T-1,K
    hist = p[:, :-1]
    # bounded finite test family: constant, each current coordinate, and tanh coordinates
    feats = [torch.ones_like(hist[..., :1]), hist, torch.tanh(3.0 * (hist - 1/3))]
    feat = torch.cat(feats, dim=-1)           # B,T-1,R
    moments = torch.einsum('btk,btr->tkr', d, feat) / p.shape[0]
    return (moments.square().mean(), moments.detach())


def ece(probs, y, bins=10):
    conf, pred = probs.max(-1)
    out = 0.0
    for lo, hi in zip(torch.linspace(0, 1, bins+1)[:-1], torch.linspace(0, 1, bins+1)[1:]):
        mask = (conf >= lo) & ((conf < hi) if hi < 1 else (conf <= hi))
        if mask.any(): out += mask.float().mean().item() * (conf[mask].mean().item() - (pred[mask] == y[mask]).float().mean().item())**2
    return math.sqrt(out)


def evaluate(model, x, y):
    model.eval()
    with torch.no_grad():
        logits = model(x); p = logits.softmax(-1)
        nll = F.cross_entropy(logits[:, -1], y).item()
        acc = (logits[:, -1].argmax(-1) == y).float().mean().item()
        step_nll = F.cross_entropy(logits.reshape(-1, 3), y[:, None].expand(-1, logits.shape[1]).reshape(-1)).item()
        step_ece = [ece(p[:, i], y) for i in range(p.shape[1])]
        ml, moments = mart_loss(logits)
        # Per-step held-out moment norm, with the same finite feature family.
        norms = moments.square().sum(dim=(1,2)).sqrt().cpu().tolist()
    return {"final_nll": nll, "final_accuracy": acc, "all_step_nll": step_nll,
            "final_ece": step_ece[-1], "step_ece": step_ece,
            "heldout_moment_norm_mean": float(np.mean(norms)),
            "heldout_moment_norm_last_transition": float(norms[-1])}


def train(lam, xtr, ytr, xte, yte, seed=SEED):
    seed_all(seed)
    model = SeqClassifier()
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    # Full-batch is deliberately small and makes the empirical moment definition transparent.
    for epoch in range(90):
        model.train(); opt.zero_grad()
        logits = model(xtr)
        loss = F.cross_entropy(logits.reshape(-1, 3), ytr[:, None].expand(-1, 6).reshape(-1))
        ml, _ = mart_loss(logits)
        (loss + lam * ml).backward(); opt.step()
    return evaluate(model, xte, yte)


def main():
    seed_all()
    toy = toy_check()
    xtr, ytr = make_data(768, seed=SEED)
    xte, yte = make_data(768, seed=SEED+1)
    results = {"toy_check": toy, "runs": {}}
    for lam in (0.0, 0.1, 1.0):
        results["runs"][str(lam)] = train(lam, xtr, ytr, xte, yte)
    results["comparison"] = {
        "baseline": results["runs"]["0.0"],
        "idea_lambda_0.1": results["runs"]["0.1"],
        "idea_lambda_1.0": results["runs"]["1.0"]}
    Path("results.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == "__main__": main()