"""Reproduce Article 02 teaching figures; no trained-model results.

Run from any working directory with Python 3.10+, NumPy, SciPy, and Matplotlib.
SVG and PNG outputs go to assets/img/posts/first-principles-02.
All logarithms are natural. Each figure has its own Axes.
"""
from __future__ import annotations

import json
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray
from scipy.integrate import quad
from scipy.optimize import minimize
from scipy.special import entr, logsumexp

Array = NDArray[np.float64]
ROOT = Path(__file__).resolve().parent
OUT = ROOT.parent / "img" / "posts" / "first-principles-02"
OUT.mkdir(parents=True, exist_ok=True)


def save(fig: plt.Figure, stem: str) -> None:
    """Save a vector master and a raster preview, without overriding colors."""
    fig.savefig(OUT / f"{stem}.svg", bbox_inches="tight")
    fig.savefig(OUT / f"{stem}.png", dpi=160, bbox_inches="tight")
    plt.close(fig)


def log_normal(x: Array, mean: float, std: float) -> Array:
    if std <= 0:
        raise ValueError("Standard deviation must be positive.")
    return -0.5 * ((x - mean) / std) ** 2 - np.log(std) - 0.5 * np.log(2 * np.pi)


def log_target(x: Array) -> Array:
    """Equal-weight mixture of two well-separated normal densities."""
    return logsumexp(
        np.stack([log_normal(x, -3.0, 0.6), log_normal(x, 3.0, 0.6)]),
        axis=0,
    ) - np.log(2.0)


def main() -> None:
    p_data = np.array([0.6, 0.2, 0.2], dtype=np.float64)
    z = np.array([0.0, 1.0, 2.0], dtype=np.float64)
    log_q = z - logsumexp(z)
    q_model = np.exp(log_q)
    h_parts = -p_data * np.log(p_data)
    ce_parts = -p_data * log_q
    h = float(h_parts.sum())
    ce = float(ce_parts.sum())
    kl = ce - h
    reverse_kl = float(np.sum(q_model * (log_q - np.log(p_data))))
    assert np.isclose(kl, np.sum(p_data * (np.log(p_data) - log_q)))

    # Figure 1: H depends on a distribution, not on a single observed outcome.
    p = np.linspace(0.0, 1.0, 1001)
    h_binary = entr(p) + entr(1.0 - p)
    fig, ax = plt.subplots(figsize=(8.3, 4.7), layout="constrained")
    ax.plot(p, h_binary, linewidth=2.3)
    ax.scatter([0, 0.5, 1], [0, np.log(2), 0], s=32)
    ax.annotate("Maximum: ln(2) = 0.6931 nats = 1 bit",
                xy=(0.5, np.log(2)), xytext=(0.5, 0.78),
                ha="center", arrowprops={"arrowstyle": "->"})
    ax.annotate("Almost certain is not very uncertain",
                xy=(0.01, float(entr(0.01) + entr(0.99))),
                xytext=(0.25, 0.25), arrowprops={"arrowstyle": "->"})
    ax.set(xlim=(0, 1), ylim=(-0.025, 0.87),
           xlabel=r"p, where $P_p=(p,1-p)$", ylabel="Entropy (nats)",
           title="Entropy measures uncertainty before the outcome")
    ax.grid(alpha=0.2)
    save(fig, "01-binary-entropy")
    assert h_binary[0] == 0 and h_binary[-1] == 0
    assert np.argmax(h_binary) == 500

    # Figure 2: The bar heights are weighted contributions, not raw surprises.
    xs = np.arange(3)
    fig, ax = plt.subplots(figsize=(8.3, 4.7), layout="constrained")
    bars_h = ax.bar(xs - 0.19, h_parts, width=0.38,
                    label=r"$-P_i\log P_i$ (entropy)")
    bars_ce = ax.bar(xs + 0.19, ce_parts, width=0.38, hatch="//",
                    label=r"$-P_i\log Q_i$ (cross-entropy)")
    ax.bar_label(bars_h, fmt="%.3f", padding=4)
    ax.bar_label(bars_ce, fmt="%.3f", padding=4)
    ax.set_xticks(xs, ["cats\nP = 0.6", "dogs\nP = 0.2", "pizza\nP = 0.2"])
    ax.set(ylim=(0, 1.68), ylabel="Contribution to the mean (nats)",
           title=f"Same observations, different scoring distributions\n"
                 f"H(P) = {h:.4f}    H(P,Q) = {ce:.4f}    KL(P||Q) = {kl:.4f}")
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.19), ncol=2, frameon=False)
    ax.grid(axis="y", alpha=0.2)
    save(fig, "02-weighted-contributions")

    # Figure 3: A controlled 1-D slice through the space of model distributions.
    q = np.linspace(0.005, 0.995, 1001)
    family = np.stack([q, (1.0-q)/2, (1.0-q)/2], axis=1)
    family_ce = -np.sum(p_data * np.log(family), axis=1)
    family_kl = family_ce - h
    fig, ax = plt.subplots(figsize=(8.3, 4.9), layout="constrained")
    ax.plot(q, family_ce, linewidth=2.2, label=r"$H(P,Q_q)$")
    ax.axhline(h, linestyle="--", linewidth=1.8, label=r"$H(P)$: fixed baseline")
    ax.plot(q, family_kl, linestyle="-.", linewidth=1.8, label=r"$D_{KL}(P\Vert Q_q)$")
    ax.scatter([0.6, 0.6], [h, 0], s=35)
    ax.annotate("q = 0.6: Q = P", xy=(0.6, h), xytext=(0.43, 2.13),
                arrowprops={"arrowstyle": "->"})
    chosen_q = 0.1
    gap_ce = -0.6*np.log(chosen_q) - 0.4*np.log((1.0-chosen_q)/2)
    ax.annotate("", xy=(chosen_q, gap_ce), xytext=(chosen_q, h),
                arrowprops={"arrowstyle": "<->"})
    ax.text(0.115, h + 0.17, "gap = KL", va="center")
    ax.set(xlim=(0, 1), ylim=(-0.06, 3.6), ylabel="Average log-loss / divergence (nats)",
           xlabel=r"q, where $Q_q=(q,(1-q)/2,(1-q)/2)$",
           title="Keep P = (0.6, 0.2, 0.2) fixed; change the model Q")
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18), ncol=2, frameon=False)
    ax.grid(alpha=0.2)
    save(fig, "03-cross-entropy-gap")
    assert np.all(family_kl >= -1e-12)
    assert np.isclose(-np.sum(p_data*np.log([0.6,0.2,0.2])), h)

    # Figure 4 (optional): restricted approximation family, not a general theorem.
    # Forward KL to a Gaussian has the exact moment-matching optimum.
    forward_mean = 0.0
    forward_std = float(np.sqrt(3.0**2 + 0.6**2))
    nodes, weights = np.polynomial.hermite.hermgauss(160)
    weights = weights / np.sqrt(np.pi)

    def reverse_objective(params: Array) -> float:
        mean, log_std = params
        std = float(np.exp(log_std))
        points = mean + np.sqrt(2.0)*std*nodes
        expected_log_q = -log_std - 0.5*np.log(2*np.pi) - 0.5
        return float(expected_log_q - np.sum(weights*log_target(points)))

    starts = [(0, np.log(3.0)), (3.0, np.log(0.6)), (-3.0, np.log(0.6)),
              (1.5, np.log(1.5)), (-1.5, np.log(1.5))]
    fits = [minimize(reverse_objective, np.asarray(start), method="L-BFGS-B",
                     bounds=[(-6, 6), (np.log(0.1), np.log(6))])
            for start in starts]
    good = [fit for fit in fits if fit.success]
    if not good:
        raise RuntimeError("All illustrative reverse-KL optimizations failed.")
    best = min(good, key=lambda fit: fit.fun)
    reverse_mean = abs(float(best.x[0]))  # symmetry gives an equally good left mode
    reverse_std = float(np.exp(best.x[1]))

    def reverse_integrand(t: float) -> float:
        xx = reverse_mean + reverse_std*t
        lq = float(log_normal(np.asarray(xx), reverse_mean, reverse_std))
        lp = float(log_target(np.asarray(xx)))
        return float(np.exp(-t*t/2)/np.sqrt(2*np.pi) * (lq-lp))

    quadrature, error = quad(reverse_integrand, -12, 12, epsabs=1e-10, limit=200)
    assert abs(quadrature-best.fun) < 1e-6
    x = np.linspace(-8, 8, 2401)
    fig, ax = plt.subplots(figsize=(8.3, 4.9), layout="constrained")
    ax.plot(x, np.exp(log_target(x)), linewidth=2.4, label="Target P: two peaks")
    ax.plot(x, np.exp(log_normal(x, forward_mean, forward_std)), linestyle="--",
            linewidth=2, label="Forward KL fit: cover both peaks")
    ax.plot(x, np.exp(log_normal(x, reverse_mean, reverse_std)), linestyle="-.",
            linewidth=2, label="Reverse KL fit: one symmetric solution")
    ax.set(xlim=(-8, 8), ylim=(0, 0.74), xlabel="x", ylabel="Probability density",
           title="One Gaussian Q cannot reproduce two separated peaks")
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.17), frameon=False)
    ax.grid(alpha=0.2)
    save(fig, "04-restricted-kl-fits")

    checks = {
        "scope": "illustrative calculations, not training experiments",
        "log_base": "e",
        "P": p_data.tolist(), "Q": q_model.tolist(),
        "entropy_nats": h, "cross_entropy_nats": ce,
        "forward_KL_nats": kl, "reverse_KL_nats": reverse_kl,
        "weighted_entropy_contributions": h_parts.tolist(),
        "weighted_CE_contributions": ce_parts.tolist(),
        "signed_KL_contributions": (ce_parts-h_parts).tolist(),
        "binary_entropy_0.01_nats": float(entr(.01)+entr(.99)),
        "information_rare_0.01_nats": float(-np.log(.01)),
        "binary_max_nats": float(np.log(2)),
        "optional_gaussian_figure": {
            "target_means": [-3, 3], "target_stds": [.6, .6], "target_weights": [.5, .5],
            "forward_fit_exact": {"mean": forward_mean, "std": forward_std},
            "reverse_fit_numerical": {"mean": reverse_mean, "std": reverse_std,
                                       "KL": float(best.fun), "quad_check": quadrature,
                                       "quad_estimated_error": error},
            "limitation": "Numerical reverse-KL fit checked from multiple starts; not a proof of global optimality."
        }
    }
    (OUT / "numerical_checks.json").write_text(json.dumps(checks, indent=2), encoding="utf-8")
    print("Created four independent figure previews and a numerical-check record.")
    print(f"H={h:.4f}, CE={ce:.4f}, forward KL={kl:.4f}, reverse KL={reverse_kl:.4f} nats")


if __name__ == "__main__":
    main()