# /// script
# requires-python = "==3.12.*"
# dependencies = [
#   "numpy",
#   "matplotlib",
#   "scipy",
# ]
# ///
"""
Generate Gamma posterior evolution figure for the channel-hopping post.

Shows Gamma posteriors over Poisson rates (conjugate prior for Poisson likelihood)
for 3 channels at t=0, t=50, t=500 under Thompson Sampling.
"""

import sys
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from scipy import stats

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from figure_quality import configure_matplotlib, figure_output_path, validate_figure

configure_matplotlib()

RNG = np.random.default_rng(42)

TRUE_RATES = [45, 30, 10]  # lambda: high, medium, low traffic (packets/s)
DWELL_SECONDS = 0.1
COLORS = ["#2171b5", "#c05000", "#737373"]  # blue, orange, gray
LABELS = [
    r"Channel 1 ($\lambda=45$)",
    r"Channel 2 ($\lambda=30$)",
    r"Channel 3 ($\lambda=10$)",
]

# Gamma(alpha, 1/beta) is conjugate prior for Poisson(lambda).
# After exposure T and total count s: posterior is Gamma(alpha + s, 1/(beta + T)).
# We parameterize as (alpha, beta) where mean = alpha/beta.
PRIOR_ALPHA = 1.0
PRIOR_BETA = 0.1  # mean 10 packets/s, broad on the scale used here


def simulate_thompson(
    checkpoints: tuple[int, ...],
) -> dict[int, list[tuple[float, float]]]:
    """
    Run one Thompson-sampling trajectory and retain posterior snapshots.
    """
    alphas = [PRIOR_ALPHA] * 3
    betas = [PRIOR_BETA] * 3

    snapshots: dict[int, list[tuple[float, float]]] = {}
    for step in range(1, max(checkpoints) + 1):
        # Sample from each posterior Gamma(alpha, 1/beta), mean = alpha/beta
        samples = [RNG.gamma(a, 1.0 / b) for a, b in zip(alphas, betas)]
        chosen = int(np.argmax(samples))
        # Observe Poisson draw from chosen channel
        obs = RNG.poisson(TRUE_RATES[chosen] * DWELL_SECONDS)
        alphas[chosen] += obs
        betas[chosen] += DWELL_SECONDS
        if step in checkpoints:
            snapshots[step] = list(zip(alphas, betas, strict=True))

    return snapshots


def gamma_pdf_range(alpha: float, beta: float, n_pts: int = 500):
    """Return x, pdf arrays for Gamma(alpha, 1/beta) covering the meaningful support."""
    dist = stats.gamma(a=alpha, scale=1.0 / beta)
    lo, hi = dist.ppf(0.001), dist.ppf(0.999)
    # Pad slightly and clip to non-negative
    lo = max(0.0, lo - (hi - lo) * 0.05)
    hi = hi + (hi - lo) * 0.05
    x = np.linspace(lo, hi, n_pts)
    return x, dist.pdf(x)


def make_panel(
    ax, posteriors: list[tuple[float, float]], title: str, show_ylabel: bool
):
    for i, ((alpha, beta), color, label) in enumerate(zip(posteriors, COLORS, LABELS)):
        x, y = gamma_pdf_range(alpha, beta)
        ax.plot(x, y, color=color, linewidth=1.8, label=label)
        ax.fill_between(x, y, alpha=0.25, color=color)

    ax.set_title(title, fontsize=13, pad=8)
    ax.set_xlabel(r"Rate $\lambda$", fontsize=12)
    if show_ylabel:
        ax.set_ylabel("Density", fontsize=12)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_linewidth(0.8)
    ax.spines["bottom"].set_linewidth(0.8)
    ax.tick_params(labelsize=10, length=3)
    ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda v, _: f"{v:.2g}"))
    ax.set_xlim(left=0)
    ax.set_ylim(bottom=0)


def main():
    prior = [(PRIOR_ALPHA, PRIOR_BETA)] * 3
    snapshots = simulate_thompson((50, 500))
    post_50 = snapshots[50]
    post_500 = snapshots[500]

    # 680px wide at 96 dpi ≈ 7.08in; use 7.2in to stay just under 680px at 96dpi
    fig, axes = plt.subplots(1, 3, figsize=(10, 3.2))
    fig.patch.set_facecolor("white")

    panels = [
        (prior, "Prior ($t=0$)", True),
        (post_50, "After 50 dwells", False),
        (post_500, "After 500 dwells", False),
    ]

    for ax, (posteriors, title, show_y) in zip(axes, panels):
        make_panel(ax, posteriors, title, show_y)

    # Shared legend below all panels
    handles, labels_ = axes[0].get_legend_handles_labels()
    fig.legend(
        handles,
        labels_,
        loc="lower center",
        ncol=3,
        fontsize=11,
        frameon=False,
        bbox_to_anchor=(0.5, -0.04),
    )

    fig.tight_layout(rect=[0, 0.08, 1, 1])

    validate_figure(fig)
    out = figure_output_path(__file__, "thompson_posteriors.png")
    fig.savefig(out, dpi=300, bbox_inches="tight", facecolor="white")
    print(f"Saved {out}")


if __name__ == "__main__":
    main()