import numpy as np
import torch
from torch.utils.data import DataLoader

from giant.constants import LOCAL_TARGET_NAMES
from giant.sample import resolve_n_sec, sample_stage1, sample_stage2

_SEC_PHYS_NAMES = ["log_mass", "charge"]


def _histogram_kl(p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8) -> float:
    """KL(P || Q) between two 1D samples, estimated via a shared histogram."""
    lo = min(p_samples.min(), q_samples.min())
    hi = max(p_samples.max(), q_samples.max())
    if hi <= lo:
        return 0.0
    edges = np.linspace(lo, hi, bins + 1)
    p_hist, _ = np.histogram(p_samples, bins=edges)
    q_hist, _ = np.histogram(q_samples, bins=edges)
    p = p_hist.astype(np.float64) + eps
    q = q_hist.astype(np.float64) + eps
    p /= p.sum()
    q /= q.sum()
    return float(np.sum(p * np.log(p / q)))


def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
    """Fraction of samples per integer value in [0, minlength), as a distribution."""
    counts = np.bincount(x, minlength=minlength)[:minlength].astype(np.float64)
    total = counts.sum()
    return counts / total if total > 0 else counts


def _categorical_kl(real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8) -> float:
    """KL(P_real || Q_gen) between two class-index samples over `n_classes`
    categories, estimated from bincount fractions. NaN if either side has no
    valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
    if len(real_idx) == 0 or len(gen_idx) == 0:
        return float("nan")
    p = _bincount_frac(real_idx, n_classes) + eps
    q = _bincount_frac(gen_idx, n_classes) + eps
    p /= p.sum()
    q /= q.sum()
    return float(np.sum(p * np.log(p / q)))


def _embedding_nearest_class(vectors: torch.Tensor, emb_weight: torch.Tensor) -> np.ndarray:
    """Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
    (vocab, emb_dim) — same computation as
    `giant.particles.decode_embedding_nearest`, but returning the raw class
    index instead of a decoded PDG code: validate.py only needs a
    real-vs-generated class-distribution comparison, not a rollout-usable
    identity, so there's no need for the pdg_map inversion here."""
    flat = vectors.reshape(-1, vectors.size(-1))
    dist = (flat.unsqueeze(1) - emb_weight.detach().unsqueeze(0)).abs().sum(-1)
    nearest = dist.argmin(dim=1)
    return nearest.reshape(vectors.shape[:-1]).cpu().numpy()


def validate_marginals(
    stage1_model: torch.nn.Module,
    val_loader: DataLoader,
    device: torch.device | None = None,
    n_batches: int | None = None,
    kl_bins: int = 50,
    steps: int = 10,
    ddpm_steps: int = 1000,
    sec_decoder: torch.nn.Module | None = None,
) -> dict[str, np.ndarray | float]:
    """Compare per-dimension marginals of generated vs. real steps.

    Returns {"real": (N,9), "generated": (N,9), "kl_divergence": (9,)} in
    normalised space. `kl_divergence[j]` is KL(real || generated) for
    dimension j, estimated from a shared histogram over both samples.

    Stage 1 is sampled via `giant.sample.sample_stage1`, which dispatches on
    `stage1_model.generator_kind` — `steps`/`ddpm_steps` are forwarded but
    only one of them is actually read, depending on that dispatch.

    When `sec_decoder` is given, also validates Stage 2 via
    `giant.sample.sample_stage2`/`resolve_n_sec` (generator- and
    one-shot-vs-autoregressive-agnostic): n_sec
    distribution (+ classification accuracy), per-slot energy-fraction
    marginals, and a particle-type marginal whose shape depends on
    `sec_decoder.particle_type_cfg.target` — restricted to each side's own
    valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since
    the two need not agree on how many slots are valid. Adds {"n_sec_real",
    "n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under
    `target = "physical"`, {"phys_real", "phys_generated", "phys_kl"}
    (continuous log_mass/charge marginals — v0.2 behaviour), or under
    `target` in `("onehot", "embedding")`, {"type_class_real",
    "type_class_gen", "type_class_kl"} (categorical class-index marginal:
    argmax for "onehot", L1-nearest conditioning-embedding row for
    "embedding" — see `_embedding_nearest_class`). Compared directly in
    normalised space (no denormalising — KL estimated from a shared
    per-sample histogram/bincount is invariant to a shared affine rescaling
    of both sides).
    """
    if device is None:
        device = next(stage1_model.parameters()).device
    stage1_model.eval()
    if sec_decoder is not None:
        sec_decoder.eval()

    k_max = sec_decoder.k_max if sec_decoder is not None else 0
    target = sec_decoder.particle_type_cfg.target if sec_decoder is not None else "physical"

    all_real, all_gen = [], []
    all_n_sec_real, all_n_sec_pred = [], []
    all_phys_real, all_phys_gen = [], []
    all_type_class_real, all_type_class_gen = [], []
    all_frac_real: list[list[np.ndarray]] = [[] for _ in range(k_max)]
    all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(k_max)]

    for i, batch in enumerate(val_loader):
        if n_batches is not None and i >= n_batches:
            break
        # batch is a StepBatch (giant.data.dataset).
        x1, n_sec, sec_cont, sec_type_idx = batch.target_s1, batch.n_sec, batch.sec_cont, batch.sec_type_idx
        cond_cont = batch.cond_cont.to(device)
        cond_cat = batch.cond_cat.to(device)

        gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)

        all_real.append(x1.numpy())
        all_gen.append(gen.cpu().numpy())

        if sec_decoder is None:
            continue

        n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
        n_sec_np = n_sec.numpy()

        real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None]  # (B, k_max)
        real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))

        sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
            sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
        )
        # A stop-token decoder resolves n_sec_pred=None above — read the real
        # count back off sec_valid_pred instead (a no-op round trip under
        # every other n_sec.mode, where sec_valid_pred was built FROM
        # n_sec_pred in the first place).
        n_sec_pred_np = sec_valid_pred.sum(dim=-1).cpu().numpy()
        all_n_sec_real.append(n_sec_np)
        all_n_sec_pred.append(n_sec_pred_np)
        gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
        gen_valid = sec_valid_pred.cpu().numpy()

        if target == "physical":
            real_phys = sec_cont[:, :, 4:6].numpy()  # (B, k_max, 2) [log_mass, charge]
            gen_phys = sec_type_pred.cpu().numpy()
            all_phys_real.append(real_phys[real_valid])
            all_phys_gen.append(gen_phys[gen_valid])
        else:
            sec_type_idx_np = sec_type_idx.numpy()
            all_type_class_real.append(sec_type_idx_np[real_valid])
            if target == "onehot":
                gen_class = sec_type_pred.argmax(dim=-1).cpu().numpy()
            else:  # "embedding"
                emb_weight = sec_decoder.cond_enc.pdg_emb.weight
                gen_class = _embedding_nearest_class(sec_type_pred, emb_weight)
            all_type_class_gen.append(gen_class[gen_valid])

        for j in range(k_max):
            all_frac_real[j].append(real_frac[real_valid[:, j], j])
            all_frac_gen[j].append(gen_frac[gen_valid[:, j], j])

    real = np.concatenate(all_real, axis=0)
    generated = np.concatenate(all_gen, axis=0)

    kl_divergence = np.array([_histogram_kl(real[:, j], generated[:, j], bins=kl_bins) for j in range(real.shape[1])])

    header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
    print(f"\n{header}")
    print("-" * len(header))
    for j, name in enumerate(LOCAL_TARGET_NAMES):
        r, g = real[:, j], generated[:, j]
        print(
            f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
            f"{r.std():>10.4f} {g.std():>10.4f} {kl_divergence[j]:>14.4f}"
        )

    result: dict[str, np.ndarray | float] = {
        "real": real,
        "generated": generated,
        "kl_divergence": kl_divergence,
    }
    if sec_decoder is None:
        return result

    n_sec_real = np.concatenate(all_n_sec_real, axis=0)
    n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0)
    n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())

    energy_fraction_kl = np.full(k_max, np.nan)
    print(f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}")
    n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
    print(n_sec_dist_header)
    print("-" * len(n_sec_dist_header))
    max_n_sec = max(int(n_sec_real.max()), int(n_sec_pred_all.max())) + 1
    real_n_sec_dist = _bincount_frac(n_sec_real, max_n_sec)
    gen_n_sec_dist = _bincount_frac(n_sec_pred_all, max_n_sec)
    for v in range(max_n_sec):
        print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}")

    print(
        f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} "
        f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
    )
    print("-" * 90)
    for j in range(k_max):
        r = np.concatenate(all_frac_real[j]) if all_frac_real[j] else np.array([])
        g = np.concatenate(all_frac_gen[j]) if all_frac_gen[j] else np.array([])
        if len(r) == 0 or len(g) == 0:
            continue
        kl = _histogram_kl(r, g, bins=kl_bins)
        energy_fraction_kl[j] = kl
        print(f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}")

    result.update(
        {
            "n_sec_real": n_sec_real,
            "n_sec_pred": n_sec_pred_all,
            "n_sec_accuracy": n_sec_accuracy,
            "energy_fraction_kl": energy_fraction_kl,
        }
    )

    if target == "physical":
        phys_real = np.concatenate(all_phys_real, axis=0)  # (M, 2)
        phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)

        if len(phys_real) > 0 and len(phys_gen) > 0:
            phys_kl = np.array([_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)])
        else:
            phys_kl = np.full(2, np.nan)

        print(
            f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
            f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
        )
        print("-" * 68)
        for j, name in enumerate(_SEC_PHYS_NAMES):
            r, g = phys_real[:, j], phys_gen[:, j]
            if len(r) == 0 or len(g) == 0:
                continue
            print(
                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
            )

        result.update({"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl})
    else:
        type_class_real = np.concatenate(all_type_class_real, axis=0)
        type_class_gen = np.concatenate(all_type_class_gen, axis=0)
        n_classes = sec_decoder.type_dim if target == "onehot" else sec_decoder.cond_enc.pdg_emb.weight.size(0)
        type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)

        print(
            f"\n{'sec type class (' + target + ')':<24} "
            f"n={len(type_class_real)}/{len(type_class_gen)} "
            f"KL(real||gen)={type_class_kl:.4f}"
        )

        result.update(
            {
                "type_class_real": type_class_real,
                "type_class_gen": type_class_gen,
                "type_class_kl": type_class_kl,
            }
        )

    return result