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

from giant.constants import LOCAL_TARGET_NAMES
from giant.sample import sample_flow, sample_ddpm, sample_ddim


def _kw(steps: int | None) -> dict[str, int]:
    return {} if steps is None else {"steps": steps}


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 validate_marginals(
    model: torch.nn.Module,
    val_loader: DataLoader,
    mode: str = "flow",
    schedule=None,
    device: torch.device | None = None,
    n_batches: int | None = None,
    kl_bins: int = 50,
    steps: int | None = None,
) -> dict[str, np.ndarray]:
    """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.

    `steps` overrides the number of sampler steps (flow ODE steps or DDIM
    substeps); `None` keeps each sampler's own default. Unused in "ddpm"
    mode, which always runs the full schedule.
    """
    if device is None:
        device = next(model.parameters()).device
    model.eval()

    all_real, all_gen = [], []
    for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
        if n_batches is not None and i >= n_batches:
            break
        cond_cont = cond_cont.to(device)
        cond_cat = cond_cat.to(device)

        if mode == "flow":
            gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
        elif mode == "ddpm":
            gen = sample_ddpm(model, cond_cont, cond_cat, schedule)
        else:
            gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))

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

    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} "
        f"{'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}"
        )

    return {"real": real, "generated": generated, "kl_divergence": kl_divergence}
