import json
import random
import numpy as np
import torch
import torch.nn as nn

SEED = 236
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'


def tv(p, q):
    return 0.5 * torch.abs(p - q).sum(dim=-1)


def make_data(n, vocab=24, modes=3, hdim=3, noise=0.35, seed=236):
    g = np.random.default_rng(seed)
    # Distinct recurring continuation distributions, with modest within-mode variation.
    centers = g.dirichlet(np.full(vocab, 0.25), size=modes).astype('float32')
    labels = g.integers(0, modes, size=n)
    p = np.empty((n, vocab), dtype='float32')
    for i, z in enumerate(labels):
        concentration = 70.0 * centers[z] + 0.25
        p[i] = g.dirichlet(concentration).astype('float32')
    # Anchor representation is informative but imperfect: one-hot mode signal + Gaussian noise.
    h = np.eye(modes, hdim, dtype='float32')[labels]
    h += noise * g.normal(size=h.shape).astype('float32')
    return torch.tensor(h), torch.tensor(p), torch.tensor(labels), torch.tensor(centers)


class SharedProposal(nn.Module):
    def __init__(self, vocab):
        super().__init__()
        self.logits = nn.Parameter(torch.zeros(vocab))
    def forward(self, h):
        return self.logits.softmax(-1).expand(h.shape[0], -1)


class RoutedPrototypes(nn.Module):
    def __init__(self, hdim, vocab, k):
        super().__init__()
        self.prototypes = nn.Parameter(torch.randn(k, vocab) * 0.1)
        self.router = nn.Linear(hdim, k)
    def forward(self, h):
        q = self.prototypes.softmax(-1)
        pi = self.router(h).softmax(-1)
        mixture = pi @ q
        return mixture, q, pi


def train_shared(h, p, steps=900, lr=0.08):
    m = SharedProposal(p.shape[1]).to(DEVICE)
    opt = torch.optim.Adam(m.parameters(), lr=lr)
    h, p = h.to(DEVICE), p.to(DEVICE)
    for _ in range(steps):
        loss = tv(p, m(h)).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return m


def train_routed(h, p, k, steps=1200, lr=0.06, entropy=0.002):
    m = RoutedPrototypes(h.shape[1], p.shape[1], k).to(DEVICE)
    opt = torch.optim.Adam(m.parameters(), lr=lr)
    h, p = h.to(DEVICE), p.to(DEVICE)
    for _ in range(steps):
        mixture, q, pi = m(h)
        # The paper's soft surrogate, with a small positive entropy penalty.
        distances = 0.5 * torch.abs(p[:, None, :] - q[None, :, :]).sum(-1)
        loss = (pi * distances).sum(-1).mean() + entropy * (-(pi * (pi + 1e-8).log()).sum(-1).mean())
        opt.zero_grad(); loss.backward(); opt.step()
    return m


def evaluate(shared, routed, h, p, labels):
    hdev, pdev = h.to(DEVICE), p.to(DEVICE)
    with torch.no_grad():
        qb = shared(hdev)
        mix, q, pi = routed(hdev)
        d = 0.5 * torch.abs(pdev[:, None, :] - q[None, :, :]).sum(-1)
        oracle = d.min(-1).values
        selected = d.gather(1, pi.argmax(-1, keepdim=True)).squeeze(1)
        actual_mix = tv(pdev, mix)
        # A simple rejection/verification proxy: maximal coupling acceptance.
        out = {
            'shared_tv': tv(pdev, qb).mean().item(),
            'oracle_prototype_tv': oracle.mean().item(),
            'routed_selected_tv': selected.mean().item(),
            'mixture_tv': actual_mix.mean().item(),
            'oracle_reduction_vs_shared': (tv(pdev, qb).mean() - oracle.mean()).item(),
            'routed_reduction_vs_shared': (tv(pdev, qb).mean() - selected.mean()).item(),
            'mixture_reduction_vs_shared': (tv(pdev, qb).mean() - actual_mix.mean()).item(),
            'shared_accept_proxy': (1-tv(pdev, qb)).mean().item(),
            'routed_accept_proxy': (1-selected).mean().item(),
            'mixture_accept_proxy': (1-actual_mix).mean().item(),
            'router_accuracy': (pi.argmax(-1).cpu() == labels).float().mean().item(),
            'prototype_count': q.shape[0],
        }
    return out


def exact_math_check():
    # Two sharply separated distributions: K=1 has a nonzero information floor,
    # while two prototypes attain zero oracle distance.
    p = torch.tensor([[0.92, 0.04, 0.02, 0.02], [0.03, 0.03, 0.04, 0.90]])
    q_shared = p.mean(0, keepdim=True).expand_as(p)
    q_proto = p.clone()
    l1 = tv(p, q_shared).mean().item()
    lk = tv(p, q_proto).min(-1).values.mean().item()
    return {'two_mode_L1': l1, 'two_mode_L2_oracle': lk, 'reduction': l1-lk,
            'bound_sanity': bool(l1 >= lk - 1e-8 and lk < 1e-8)}


def main():
    math = exact_math_check()
    htr, ptr, ytr, _ = make_data(1800, seed=SEED)
    hte, pte, yte, _ = make_data(900, seed=SEED+1)
    # Use the same underlying centers for a fair held-out test by generating test labels
    # from training centers with fresh within-mode samples.
    rng = np.random.default_rng(SEED+1)
    centers = np.array(make_data(1, seed=SEED)[3])
    labels = rng.integers(0, 3, size=900)
    ptest = np.array([rng.dirichlet(70*centers[z] + .25) for z in labels]).astype('float32')
    htest = np.eye(3, 3, dtype='float32')[labels] + .35*rng.normal(size=(900,3)).astype('float32')
    hte, pte, yte = torch.tensor(htest), torch.tensor(ptest), torch.tensor(labels)
    shared = train_shared(htr, ptr)
    results = {'device': DEVICE, 'math_check': math, 'train_n': len(htr), 'test_n': len(hte)}
    # Train on identical data and compare K=2 and K=3 to the standard single head.
    for k in (2, 3):
        routed = train_routed(htr, ptr, k)
        results[f'K{k}'] = evaluate(shared, routed, hte, pte, yte)
    # Parameter/FLOP accounting for the final vocabulary heads (shared hidden compute).
    results['head_parameter_multiplier_K3'] = 3.0
    results['note'] = 'Output-head-only multiplier; shared hidden/router compute is not included in this synthetic proxy.'
    print(json.dumps(results, indent=2, sort_keys=True))

if __name__ == '__main__':
    try:
        main()
    except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
        if DEVICE == 'cuda':
            print(json.dumps({'cuda_error': str(e), 'fallback': 'rerun with CUDA_VISIBLE_DEVICES empty'}))
            raise
        raise