"""Measurements of a batch of generated samples.

These are small, exactly computable quantities, each a diagnostic for whether a sampler learned the
structure planted in the data: the right letter frequencies, the planted motif, the right number of
modes.
"""

from __future__ import annotations

import torch

from cis6270.data import GAUSSIAN_CENTERS, MOTIF, decode_dna


# %% 01. Measurements on DNA


def gc_fraction(tokens: torch.Tensor) -> float:
    """Fraction of positions holding C or G, averaged over the batch. The training data sit at 0.577."""
    return float(((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean())


def motif_fraction(tokens: torch.Tensor, motif: str = MOTIF) -> float:
    """Fraction of sequences that contain the planted motif anywhere."""
    sequences = decode_dna(tokens)
    return sum(motif in sequence for sequence in sequences) / max(len(sequences), 1)


def unique_fraction(tokens: torch.Tensor) -> float:
    """Fraction of generated sequences that are distinct. A low value marks a collapsed sampler.

    We compare the token ids themselves, so this works for any vocabulary.
    """
    rows = [tuple(row.tolist()) for row in tokens]  # One tuple per sequence.
    return len(set(rows)) / max(len(rows), 1)


def token_entropy(tokens: torch.Tensor, vocab_size: int) -> float:
    """Entropy in nats of the letter frequencies, at most log K, or log 4 = 1.386 at K = 4."""
    counts = torch.bincount(tokens.flatten(), minlength=vocab_size).to(torch.get_default_dtype())
    probabilities = counts / counts.sum().clamp(min=1)
    nonzero = probabilities[probabilities > 0]
    return float(-(nonzero * nonzero.log()).sum())


def support_fraction(tokens: torch.Tensor, training: torch.Tensor) -> float:
    """Fraction of generated sequences that also appear in the training set."""
    seen = {tuple(row.tolist()) for row in training}
    return sum(tuple(row.tolist()) in seen for row in tokens) / max(len(tokens), 1)


# %% 02. Measurements on two-dimensional samples


def mode_report(points: torch.Tensor, centers: torch.Tensor | None = None) -> dict:
    """Assign each point to its nearest mode and report the distance and the coverage.

    ``nearest_distance`` near the data's own standard deviation means the points land on the modes.
    ``mode_entropy`` near log 4 = 1.386 means all four modes are visited about equally.
    """
    centers = GAUSSIAN_CENTERS.to(points) if centers is None else centers.to(points)
    distance = torch.cdist(points, centers)  # [N, M] distance to every center.
    nearest, which = distance.min(dim=1)  # [N] distance and index of the closest mode.
    counts = torch.bincount(which, minlength=len(centers)).to(torch.get_default_dtype())
    fractions = counts / counts.sum().clamp(min=1)
    nonzero = fractions[fractions > 0]
    return {
        "nearest_distance": float(nearest.mean()),
        "modes_found": int((counts > 0).sum()),
        "mode_entropy": float(-(nonzero * nonzero.log()).sum()),
        "mode_fractions": [round(float(f), 3) for f in fractions],
    }


def mean_log_density(points: torch.Tensor) -> float:
    """Average exact log density of the four-Gaussian target at the generated points."""
    from cis6270.data import gaussian_mixture_log_density

    return float(gaussian_mixture_log_density(points).mean())


# %% 03. Measurements on a learned distribution over a small state space
#
# Several chapters work on a state space small enough to enumerate, so the exact answer is
# available. We compare an empirical histogram against that exact answer, the sharpest check in
# the course.


def total_variation(p: torch.Tensor, q: torch.Tensor) -> float:
    """Half the sum of absolute differences between two distributions over the same states."""
    return float(0.5 * (p - q).abs().sum())


def empirical_distribution(tokens: torch.Tensor, vocab_size: int, length: int) -> torch.Tensor:
    """Histogram of whole sequences over the K**L states, for a state space small enough to list."""
    weights = vocab_size ** torch.arange(length - 1, -1, -1, device=tokens.device)  # Place values.
    indices = (tokens * weights).sum(dim=1)  # [N] index of each sequence in the full state space.
    counts = torch.bincount(indices, minlength=vocab_size ** length).to(torch.get_default_dtype())
    return counts / counts.sum().clamp(min=1)  # [K**L] empirical distribution.