"""Entropy, cross-entropy, KL divergence and maximum likelihood, in nats.

Information content measures one observation, and entropy averages it under the distribution that
produced it. Cross-entropy averages the same quantity under a second distribution supplying the
probabilities, and the KL divergence is the difference between those two averages. So when we
minimize a negative log-likelihood over a model, we are minimizing a KL divergence to the data
distribution.

Chapter 1, Sections 1.33, 1.34 and 1.35 of the notes. Runs on the pair p = (0.8, 0.2) and
q = (0.6, 0.4) of Section 1.33.3 and on the ten DNA bases of Section 1.35. Prints the chapter's
number beside the computed one. Logarithms are natural, so every quantity is in nats.

    python lecture_1/information.py                   # every worked example
    python lecture_1/information.py --samples 200000  # a tighter Monte Carlo cross-entropy
    python lecture_1/information.py --quiet           # the closing report only
"""

# %% 01. Imports and configuration
from __future__ import annotations

import argparse
import math

import torch

from cis6270.runner import print_report, save_run, set_seed

torch.set_default_dtype(torch.float64)  # Double precision, because the chapter prints five decimal places.


P_REF = torch.tensor([0.8, 0.2])  # The reference distribution p of Sections 1.33.3 and 1.34.
Q_MODEL = torch.tensor([0.6, 0.4])  # The comparison distribution q scored against those samples.
BASE_COUNTS = torch.tensor([4.0, 3.0, 2.0, 1.0])  # n(A), n(C), n(G), n(T) over ten DNA bases.


WORKED: dict[str, float] = {}  # Every computed value, by label, for tests/test_lecture_1.py.


def show(label: str, by_hand: float, computed: float, quiet: bool = False) -> None:
    """Print the chapter's hand-worked value beside the value this file computes, and keep it."""
    WORKED[label] = computed  # The test suite asserts these against the chapter's own numbers.
    if not quiet:  # Under --quiet only the closing report is printed.
        computed = computed + 0.0  # Adding zero prints a negative zero as a plain zero.
        print(f"  {label:<36s} chapter {by_hand:>11.5f}   computed {computed:>11.5f}"
              f"   gap {abs(by_hand - computed):.1e}")


def entropy(p: torch.Tensor) -> float:
    """H(p) = -sum_x p(x) log p(x), the expected information in one sample drawn from p."""
    return float(-(p * p.clamp_min(1e-300).log()).sum())  # Clamp keeps 0 log 0 at zero.


def cross_entropy(p: torch.Tensor, q: torch.Tensor) -> float:
    """H(p, q) = -sum_x p(x) log q(x), the cost of scoring samples from p with the values of q."""
    return float(-(p * q.log()).sum())  # The weights come from p and the probabilities from q.


def divergence(p: torch.Tensor, q: torch.Tensor) -> float:
    """KL(p || q) = sum_x p(x) log(p(x) / q(x)), the expected log ratio taken under p."""
    return float((p * (p / q).log()).sum())  # Zero exactly when q matches p at every state.


# %% 02. Information content
def information(quiet: bool = False) -> dict:
    """Section 1.33.1: I(x) = -log p(x), so a rarer outcome has more information."""
    common = -math.log(0.5)  # An outcome of probability one half.
    rare = -math.log(0.01)  # An outcome fifty times less likely.
    certain = -math.log(1.0)  # A certain event, with zero information.
    show("I(x) at p(x) = 0.5", 0.693, common, quiet)
    show("I(x) at p(x) = 0.01", 4.605, rare, quiet)
    show("I(x) at p(x) = 1", 0.0, certain, quiet)
    show("ratio of the two surprises", 6.6, rare / common, quiet)
    return {"common": common, "rare": rare}


# %% 03. Entropy
def entropies(quiet: bool = False) -> dict:
    """Section 1.33.2: entropy is high when the mass is spread and zero when it is all on one state."""
    skewed = entropy(torch.tensor([0.9, 0.1]))  # 0.09482 + 0.23026 nats.
    balanced = entropy(torch.tensor([0.5, 0.5]))  # log 2, the largest value on two states.
    deterministic = entropy(torch.tensor([1.0, 0.0]))  # All mass on one state, so 1 log 1 = 0.
    uniform_four = entropy(torch.full((4,), 0.25))  # log K at K = 4, proved in Section 1.37.2.
    show("H((0.9, 0.1))", 0.325, skewed, quiet)
    show("H((0.5, 0.5)) = log 2", 0.693, balanced, quiet)
    show("H of a point mass", 0.0, deterministic, quiet)
    show("H of the uniform on 4 states", 1.386, uniform_four, quiet)
    return {"skewed": skewed, "balanced": balanced, "uniform_four": uniform_four}


# %% 04. Cross-entropy
def cross_entropies(quiet: bool = False) -> dict:
    """Section 1.33.3: scoring samples from p with the probabilities of q costs more than H(p)."""
    cost = cross_entropy(P_REF, Q_MODEL)  # 0.40866 + 0.18326 nats.
    own = entropy(P_REF)  # 0.17851 + 0.32189 nats, the irreducible part of that cost.
    same = cross_entropy(P_REF, P_REF)  # H(p, p) recovers the entropy exactly.
    show("H(p, q)", 0.592, cost, quiet)
    show("H(p)", 0.500, own, quiet)
    show("excess cost H(p, q) - H(p)", 0.092, cost - own, quiet)
    show("H(p, p) - H(p)", 0.0, same - own, quiet)
    return {"cost": cost, "own": own, "excess": cost - own}


# %% 05. Kullback-Leibler divergence
def kl(samples: int, quiet: bool = False) -> dict:
    """Section 1.34: the split H(p, q) = H(p) + KL(p || q), nonnegativity, and the asymmetry."""
    forward = divergence(P_REF, Q_MODEL)  # 0.23015 - 0.13863 nats, the price of using q.
    reverse = divergence(Q_MODEL, P_REF)  # -0.17261 + 0.27726 nats, the other direction.
    split = entropy(P_REF) + forward  # The boxed identity, and it must return H(p, q).
    logits = torch.randn(max(samples, 1), 2, 8)  # Random logits for pairs of distributions.
    pairs = torch.softmax(logits, dim=-1)  # Two distributions over eight states per row; [N, 2, 8].
    drawn = (pairs[:, 0] * (pairs[:, 0] / pairs[:, 1]).log()).sum(-1)  # KL for every pair; [N].
    violation = float((-drawn).clamp_min(0.0).max())  # The largest breach of KL >= 0 among them.
    t = torch.linspace(0.01, 5.0, 500)  # The range on which we check the inequality log t <= t - 1.
    show("KL(p || q)", 0.0915, forward, quiet)
    show("KL(q || p)", 0.1046, reverse, quiet)
    show("H(p) + KL(p || q) = H(p, q)", 0.592, split, quiet)
    show("KL(p || p)", 0.0, divergence(P_REF, P_REF), quiet)
    show("reverse over forward, percent", 14.0, 100.0 * (reverse / forward - 1.0), quiet)
    show("largest log t - (t - 1)", 0.0, float((t.log() - (t - 1.0)).max()), quiet)
    show(f"largest KL breach over {len(drawn)} pairs", 0.0, violation, quiet)
    return {"forward": forward, "reverse": reverse, "split": split,
            "smallest": float(drawn.min()), "violation": violation}


# %% 06. Cross-entropy by Monte Carlo
def monte_carlo(samples: int, repeats: int, quiet: bool = False) -> dict:
    """Section 1.5.3 applied to Section 1.33.3: H(p, q) is an expectation, so an average estimates it.

    We draw states from p and average -log q at the drawn states. The estimate moves with --seed
    and with --samples, while the exact value it approaches stays fixed. We repeat a ten-sample
    estimate ``repeats`` times. The spread of those estimates is the sd that Section 1.5.3
    predicts for a ten-sample average.
    """
    exact = cross_entropy(P_REF, Q_MODEL)  # The target of the estimator, 0.592 nats.
    draws = torch.multinomial(P_REF, max(samples, 10), replacement=True)  # [N] states drawn from p.
    surprise = -Q_MODEL.log()  # The two values -log q(x) the estimator averages over; [2].
    scored = surprise[draws]  # The summand -log q(x^(n)) at each drawn state; [N].
    estimate, spread = float(scored.mean()), float(scored.std())  # The average and its per-sample sd.
    error = spread / math.sqrt(len(draws))  # The sd of the average, and it falls like 1/sqrt(N).
    show("H(p, q) exactly", 0.592, exact, quiet)
    if not quiet:
        label = f"H(p, q) from {len(draws)} samples"
        print(f"  {label:<36s} exact   {exact:>11.5f}   estimate {estimate:>11.5f}"
              f"   sd {error:.1e}")
    rounds = max(repeats, 2)  # Two averages is the fewest we can take a standard deviation over.
    repeated = torch.multinomial(P_REF, 10 * rounds, replacement=True)  # [10 * rounds] states.
    averages = surprise[repeated].view(rounds, 10).mean(dim=1)  # One ten-sample estimate per round.
    exact_spread = float((P_REF @ (surprise - exact).square()).sqrt())  # sd of one scored draw.
    predicted = exact_spread / math.sqrt(10.0)  # The sd of a ten-sample average, Section 1.5.3.
    show(f"sd of {rounds} ten-sample estimates", predicted, float(averages.std()), quiet)
    return {"exact": exact, "estimate": estimate, "error": error, "draws": len(draws),
            "predicted": predicted, "measured": float(averages.std()), "rounds": rounds}


# %% 07. Maximum likelihood as KL minimization
def maximum_likelihood(quiet: bool = False) -> dict:
    """Section 1.35: the empirical loss at the frequencies is their entropy, and the gap to any
    other model is the KL divergence from the frequencies to it."""
    frequencies = BASE_COUNTS / BASE_COUNTS.sum()  # p_hat = (0.4, 0.3, 0.2, 0.1) over A, C, G, T.
    uniform = torch.full((4,), 0.25)  # The categorical model that spreads the mass evenly.
    at_frequencies = cross_entropy(frequencies, frequencies)  # The empirical loss at p_hat.
    at_uniform = cross_entropy(frequencies, uniform)  # The same data scored by the uniform model.
    gap = at_uniform - at_frequencies  # The reducible part of the cost, a KL divergence.
    terms = -frequencies * frequencies.log()  # The four terms the chapter adds up; [4].
    show("term for A, 0.4 log(1/0.4)", 0.36652, float(terms[0]), quiet)
    show("term for C, 0.3 log(1/0.3)", 0.36119, float(terms[1]), quiet)
    show("term for G, 0.2 log(1/0.2)", 0.32189, float(terms[2]), quiet)
    show("term for T, 0.1 log(1/0.1)", 0.23026, float(terms[3]), quiet)
    show("loss at the frequencies", 1.280, at_frequencies, quiet)
    show("loss under the uniform model", 1.386, at_uniform, quiet)
    show("gap between the two losses", 0.1064, gap, quiet)
    show("KL(p_hat || uniform)", 0.1064, divergence(frequencies, uniform), quiet)
    return {"frequencies": frequencies, "at_frequencies": at_frequencies,
            "at_uniform": at_uniform, "gap": gap}


# %% 08. The report
def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--samples", type=int, default=20000, help="draws in the Monte Carlo checks")
    parser.add_argument("--steps", type=int, default=5000, help="repeats of the ten-sample estimate")
    parser.add_argument("--seed", type=int, default=0, help="random seed for the Monte Carlo draws")
    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    args = parser.parse_args()

    set_seed(args.seed)  # Seed first, so the draws repeat under the same seed.
    if not args.quiet:
        print("Chapter 1, Sections 1.33 to 1.35, worked example by worked example\n")
    surprise = information(args.quiet)
    spread = entropies(args.quiet)
    cost = cross_entropies(args.quiet)
    gap = kl(args.samples, args.quiet)
    estimate = monte_carlo(args.samples, args.steps, args.quiet)
    fitted = maximum_likelihood(args.quiet)

    report = {
        "I(x) at p(x) = 0.01, nats": surprise["rare"],
        "H((0.9, 0.1)), nats": spread["skewed"],
        "H((0.5, 0.5)), nats": spread["balanced"],
        "H(p), nats": cost["own"],
        "H(p, q), nats": cost["cost"],
        "KL(p || q), nats": gap["forward"],
        "KL(q || p), nats": gap["reverse"],
        "smallest KL over random pairs": gap["smallest"],
        "largest KL breach, random pairs": gap["violation"],
        "H(p, q) from samples": estimate["estimate"],
        "Monte Carlo samples drawn": estimate["draws"],
        "ten-sample estimates taken": estimate["rounds"],
        "sd of those estimates": estimate["measured"],
        "sd Section 1.5.3 predicts": estimate["predicted"],
        "empirical loss at p_hat, nats": fitted["at_frequencies"],
        "loss under the uniform, nats": fitted["at_uniform"],
        "KL(p_hat || uniform), nats": fitted["gap"],
    }
    print_report("Information, Sections 1.33 to 1.35", report)
    save_run(args.out, config=vars(args), report=report)


if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work every example of Sections 1.33 to 1.35 and print the report.