"""Create deterministic SVG figures for the probability/statistics tutorial.

Run simulate.py first, then run this script from any working directory.
The fictional counts and analytic distribution examples are documented here;
coverage values come exclusively from results/simulation.json.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
import tempfile

os.environ.setdefault("MPLCONFIGDIR", str(Path(tempfile.gettempdir()) / "statistics-mpl"))

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
from matplotlib.ticker import PercentFormatter
import numpy as np
from scipy.special import expit
from scipy.stats import binom


ROOT = Path(__file__).resolve().parents[1]
OUTPUT = ROOT / "figures"
BLUE = "#1D4E89"
GOLD = "#A85D12"
INK = "#25313D"
GREY = "#66727D"
GRID = "#DFE4E8"

plt.rcParams.update(
    {
        "font.family": "DejaVu Sans",
        "font.size": 12,
        "axes.titlesize": 14,
        "axes.labelsize": 12,
        "axes.titleweight": "bold",
        "axes.labelcolor": INK,
        "axes.edgecolor": GREY,
        "text.color": INK,
        "xtick.color": INK,
        "ytick.color": INK,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "figure.facecolor": "white",
        "axes.facecolor": "white",
        "legend.frameon": False,
        "lines.linewidth": 2,
        "svg.fonttype": "none",
        "svg.hashsalt": "probability-statistics-teaching-20261012",
    }
)


def horizontal_grid(ax: plt.Axes) -> None:
    ax.set_axisbelow(True)
    ax.grid(axis="y", color=GRID, linewidth=0.8)


def teaching_rates() -> plt.Figure:
    """Observed proportions, not estimates from real-world evidence."""
    successes = np.array([40, 70])
    sample_sizes = np.array([200, 200])
    rates = successes / sample_sizes
    fig, ax = plt.subplots(figsize=(8.5, 5.2))
    fig.subplots_adjust(left=0.13, right=0.97, top=0.82, bottom=0.19)
    fig.suptitle(
        "Success rates in the fictional experiment", y=0.96, fontsize=16, weight="bold"
    )
    fig.text(
        0.5, 0.88, "Independent individuals; random allocation", ha="center", color=GREY
    )
    bars = ax.bar([0, 1], rates, width=0.52, color=[BLUE, GOLD], zorder=3)
    ax.set_xticks([0, 1], ["Control\nn = 200", "Intervention\nn = 200"])
    ax.set_ylabel("Observed success rate")
    ax.set_ylim(0, 0.5)
    ax.set_yticks(np.arange(0, 0.51, 0.1))
    ax.yaxis.set_major_formatter(PercentFormatter(1, decimals=0))
    ax.set_xlim(-0.65, 1.65)
    horizontal_grid(ax)
    for bar, rate, count, size in zip(bars, rates, successes, sample_sizes):
        ax.text(
            bar.get_x() + bar.get_width() / 2,
            rate + 0.02,
            f"{rate:.0%}\n{count} / {size} successes",
            ha="center",
            va="bottom",
            weight="bold",
            fontsize=12,
        )
    fig.text(
        0.5, 0.035, "Teaching data only. Bars describe this sample.", ha="center", color=GREY
    )
    return fig


def distributions() -> plt.Figure:
    """The CDF is right-continuous and jumps at integers, unlike the PMF."""
    n, p = 10, 0.3
    k = np.arange(n + 1)
    mass = binom.pmf(k, n, p)
    cdf = binom.cdf(k, n, p)
    fig, axes = plt.subplots(1, 2, figsize=(10.6, 5.1))
    fig.subplots_adjust(left=0.08, right=0.97, bottom=0.16, top=0.78, wspace=0.30)
    fig.suptitle("Binomial(n = 10, p = 0.30)", y=0.96, fontsize=16, weight="bold")
    pmf_ax, cdf_ax = axes
    pmf_ax.set_title("Probability mass function", pad=14)
    pmf_ax.vlines(k, 0, mass, color=BLUE, linewidth=2)
    pmf_ax.scatter(k, mass, color=BLUE, s=38, zorder=4)
    pmf_ax.set_xlabel("Number of successes k")
    pmf_ax.set_ylabel("P(X = k)")
    pmf_ax.set_xlim(-0.6, n + 0.6)
    pmf_ax.set_ylim(0, 0.32)
    pmf_ax.set_xticks(k)
    pmf_ax.set_yticks([0, 0.1, 0.2, 0.3])
    horizontal_grid(pmf_ax)

    cdf_ax.set_title("Cumulative distribution function", pad=14)
    # F(x)=0 for x<0; for integer k, F(x)=F(k) on [k,k+1).
    step_x = np.arange(-1, n + 2)
    step_y = np.r_[0, cdf, 1]
    cdf_ax.step(step_x, step_y, where="post", color=GOLD)
    # Closed upper and open lower endpoints show the right-continuity.
    cdf_ax.scatter(k, cdf, color=GOLD, s=32, zorder=5)
    cdf_ax.scatter(
        k, np.r_[0, cdf[:-1]], facecolors="white", edgecolors=GOLD, s=32, zorder=4
    )
    cdf_ax.set_xlabel("Threshold x")
    cdf_ax.set_ylabel("P(X <= x)")
    cdf_ax.set_xlim(-0.6, n + 0.6)
    cdf_ax.set_ylim(0, 1.04)
    cdf_ax.set_xticks(k)
    cdf_ax.set_yticks(np.arange(0, 1.01, 0.2))
    horizontal_grid(cdf_ax)
    fig.text(
        0.5,
        0.025,
        "Masses sum to 1. Each CDF jump has height P(X = k).",
        ha="center",
        color=GREY,
    )
    return fig


def logistic_link() -> plt.Figure:
    """A fixed beta gives a fixed odds ratio, while the risk difference varies."""
    eta = np.linspace(-5, 5, 1001)
    beta = np.log(2)
    baseline = expit(eta)
    shifted = expit(eta + beta)
    fig, axes = plt.subplots(1, 2, figsize=(10.6, 5.5))
    fig.subplots_adjust(left=0.09, right=0.97, bottom=0.20, top=0.76, wspace=0.33)
    fig.suptitle(
        "A fixed log-odds shift and a changing risk difference",
        y=0.97,
        fontsize=15,
        weight="bold",
    )
    fig.text(
        0.5,
        0.875,
        "Holding other covariates fixed: beta = log(2), odds ratio = exp(beta) = 2",
        ha="center",
        color=GREY,
        fontsize=11,
    )
    prob_ax, rd_ax = axes
    prob_ax.set_title("Inverse logit: p = 1 / (1 + exp(-eta))", fontsize=12, pad=13)
    prob_ax.plot(eta, baseline, color=BLUE, label="Baseline eta")
    prob_ax.plot(
        eta, shifted, color=GOLD, linestyle="--", label="eta + log(2): OR = 2"
    )
    prob_ax.set_xlabel("Linear predictor eta")
    prob_ax.set_ylabel("Success probability")
    prob_ax.set_xlim(-5, 5)
    prob_ax.set_ylim(0, 1)
    prob_ax.set_xticks([-5, -2.5, 0, 2.5, 5])
    prob_ax.set_yticks(np.arange(0, 1.01, 0.2))
    prob_ax.yaxis.set_major_formatter(PercentFormatter(1, decimals=0))
    prob_ax.legend(loc="upper left", fontsize=10)
    horizontal_grid(prob_ax)

    rd_ax.set_title("Same OR = 2 at different baselines", fontsize=12, pad=13)
    rd_ax.plot(baseline, 100 * (shifted - baseline), color=GOLD)
    selected_baselines = np.array([0.1, 0.5, 0.9])
    selected_shifted = 2 * selected_baselines / (1 + selected_baselines)
    selected_rd = 100 * (selected_shifted - selected_baselines)
    rd_ax.scatter(selected_baselines, selected_rd, color=GOLD, s=42, zorder=4)
    for base, rd in zip(selected_baselines, selected_rd):
        rd_ax.annotate(
            f"{rd:.1f} pp",
            (base, rd),
            xytext=(0, 10),
            textcoords="offset points",
            ha="center",
            va="bottom",
            fontsize=10,
        )
    rd_ax.set_xlabel("Baseline success probability")
    rd_ax.set_ylabel("Risk difference (percentage points)")
    rd_ax.set_xlim(0, 1)
    rd_ax.set_ylim(0, 20)
    rd_ax.set_xticks(np.arange(0, 1.01, 0.2))
    rd_ax.xaxis.set_major_formatter(PercentFormatter(1, decimals=0))
    rd_ax.set_yticks([0, 5, 10, 15, 20])
    horizontal_grid(rd_ax)
    fig.text(
        0.5,
        0.045,
        "Risk difference = shifted probability - baseline probability. OR and RD use different scales.",
        ha="center",
        color=GREY,
        fontsize=10,
    )
    return fig


def read_coverage() -> list[dict]:
    path = ROOT / "results" / "simulation.json"
    if not path.exists():
        raise FileNotFoundError(f"Run simulate.py first: missing {path}")
    result = json.loads(path.read_text(encoding="utf-8"))
    expected_scenarios = [(20, 0.05), (100, 0.3)]
    selected = []
    for n, p in expected_scenarios:
        matches = [
            item for item in result["ci_coverage"] if item["n"] == n and item["p"] == p
        ]
        if len(matches) != 1:
            raise ValueError(f"Expected exactly one coverage result for n={n}, p={p}")
        item = matches[0]
        for method in ("wald", "wilson"):
            for key in ("exact_binomial_coverage", "empirical_coverage"):
                value = item[method][key]
                if not np.isfinite(value) or not 0 <= value <= 1:
                    raise ValueError(f"Invalid coverage: {method}.{key}={value}")
            mcse = item[method]["monte_carlo_se"]
            if not np.isfinite(mcse) or mcse < 0:
                raise ValueError(f"Invalid Monte Carlo SE: {method}={mcse}")
        selected.append(item)
    return selected


def coverage() -> plt.Figure:
    """Bars show deterministic binomial-sum coverage, not simulation estimates."""
    scenarios = read_coverage()
    fig, axes = plt.subplots(
        2,
        1,
        figsize=(9.5, 7.6),
        gridspec_kw={"height_ratios": [3.2, 1.45]},
    )
    fig.subplots_adjust(left=0.12, right=0.96, bottom=0.14, top=0.78, hspace=0.50)
    fig.suptitle("Coverage of nominal 95% proportion intervals", y=0.96, fontsize=16, weight="bold")
    fig.text(
        0.5, 0.895, "Bars: exact coverage under the stated binomial model", ha="center", color=GREY
    )
    ax, table_ax = axes
    positions = np.arange(len(scenarios))
    width = 0.30
    rows = []
    for offset, method, color in [(-width / 2, "wald", BLUE), (width / 2, "wilson", GOLD)]:
        values = [item[method]["exact_binomial_coverage"] for item in scenarios]
        bars = ax.bar(positions + offset, values, width, color=color, label=method.title())
        for bar, value in zip(bars, values):
            ax.text(
                bar.get_x() + bar.get_width() / 2,
                value - 0.045,
                f"{value:.2%}",
                ha="center",
                va="top",
                color="white",
                fontsize=11,
                weight="bold",
            )
    ax.axhline(0.95, color=INK, linestyle="--", linewidth=1.5, label="Nominal 95%")
    ax.set_xticks(positions, [f"n = {item['n']}, p = {item['p']:.2f}" for item in scenarios])
    ax.set_ylabel("Repeated-sampling coverage")
    ax.set_xlim(-0.6, 1.6)
    ax.set_ylim(0, 1)
    ax.set_yticks(np.arange(0, 1.01, 0.2))
    ax.yaxis.set_major_formatter(PercentFormatter(1, decimals=0))
    horizontal_grid(ax)
    ax.legend(loc="lower center", bbox_to_anchor=(0.5, 1.02), ncol=3, fontsize=11)
    for item in scenarios:
        for method in ("wald", "wilson"):
            values = item[method]
            rows.append(
                [
                    f"n={item['n']}, p={item['p']:.2f}",
                    method.title(),
                    f"{values['exact_binomial_coverage']:.2%}",
                    f"{values['empirical_coverage']:.2%}",
                    f"{100 * values['monte_carlo_se']:.3f} pp",
                ]
            )
    table_ax.axis("off")
    table = table_ax.table(
        cellText=rows,
        colLabels=["Scenario", "Method", "Exact coverage", "Simulation", "MCSE"],
        colWidths=[0.24, 0.16, 0.22, 0.20, 0.18],
        cellLoc="center",
        loc="center",
    )
    table.auto_set_font_size(False)
    table.set_fontsize(10)
    table.scale(1, 1.75)
    for (row, _col), cell in table.get_celld().items():
        cell.set_edgecolor(GRID)
        cell.set_linewidth(0.7)
        cell.set_text_props(color=INK)
        if row == 0:
            cell.set_facecolor("#EEF2F5")
            cell.set_text_props(weight="bold")
    repeats = {item["repeats"] for item in scenarios}
    if len(repeats) != 1:
        raise ValueError("Different repetition counts need scenario-specific figure notes")
    repeats_text = f"{next(iter(repeats)):,}"
    fig.text(
        0.5,
        0.068,
        f"Simulation: {repeats_text} repetitions per scenario. MCSE measures simulation noise.",
        ha="center",
        color=GREY,
        fontsize=10,
    )
    fig.text(
        0.5,
        0.033,
        "Exact coverage sums binomial probabilities of samples whose intervals contain p.",
        ha="center",
        color=GREY,
        fontsize=10,
    )
    return fig


def main() -> None:
    OUTPUT.mkdir(parents=True, exist_ok=True)
    factories = {
        "teaching-rates.svg": teaching_rates,
        "distributions.svg": distributions,
        "logistic-link.svg": logistic_link,
        "coverage.svg": coverage,
    }
    for name, factory in factories.items():
        figure = factory()
        target = OUTPUT / name
        figure.savefig(
            target,
            format="svg",
            bbox_inches="tight",
            pad_inches=0.16,
            metadata={"Date": None, "Creator": "Probability-Statistics tutorial"},
        )
        # Preserve SVG geometry while avoiding generated trailing whitespace.
        target.write_text(
            "\n".join(line.rstrip() for line in target.read_text(encoding="utf-8").splitlines()) + "\n",
            encoding="utf-8",
        )
        figure.savefig(
            target.with_suffix(".png"),
            format="png",
            dpi=160,
            bbox_inches="tight",
            pad_inches=0.16,
            metadata={"Software": "Probability-Statistics tutorial"},
        )
        plt.close(figure)
        print(f"Created {target}")


if __name__ == "__main__":
    main()