# /// script
# requires-python = "==3.12.*"
# dependencies = ["matplotlib", "numpy"]
# ///
"""Replicated entropy-estimator experiment for Zipf(1.1, K=5000).

The plot reports the sampling distribution, not one favorable trajectory.
Run with:

    uv run gen_convergence.py
"""

import sys
from pathlib import Path

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np

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

configure_matplotlib()

SEED = 0x51A7
SUPPORT = 5_000
ZIPF_EXPONENT = 1.1
TRIALS = 400
SAMPLE_SIZES = np.array([50, 100, 200, 400, 800, 1_600, 3_200, 6_400])


def entropy_estimates(counts: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Return plug-in, Miller--Madow, and delete-one jackknife entropy in bits."""
    sample_size = int(counts[0].sum())
    positive = counts > 0
    float_counts = counts.astype(float)
    count_log_count = np.zeros_like(float_counts)
    count_log_count[positive] = float_counts[positive] * np.log(float_counts[positive])
    sum_count_log_count = count_log_count.sum(axis=1)

    plugin_nats = np.log(sample_size) - sum_count_log_count / sample_size
    observed_support = np.count_nonzero(positive, axis=1)
    miller_madow_nats = plugin_nats + (observed_support - 1) / (2 * sample_size)

    decremented = np.maximum(float_counts - 1.0, 0.0)
    decremented_log_decremented = np.zeros_like(float_counts)
    above_one = counts > 1
    decremented_log_decremented[above_one] = decremented[above_one] * np.log(
        decremented[above_one]
    )
    removal_delta = count_log_count - decremented_log_decremented
    weighted_delta = (float_counts * removal_delta).sum(axis=1) / sample_size
    expected_leave_one_out = np.log(sample_size - 1) - (
        sum_count_log_count - weighted_delta
    ) / (sample_size - 1)
    jackknife_nats = (
        sample_size * plugin_nats - (sample_size - 1) * expected_leave_one_out
    )

    scale = 1.0 / np.log(2.0)
    return (
        plugin_nats * scale,
        miller_madow_nats * scale,
        jackknife_nats * scale,
    )


ranks = np.arange(1, SUPPORT + 1, dtype=float)
probabilities = ranks**-ZIPF_EXPONENT
probabilities /= probabilities.sum()
true_entropy = float(-(probabilities * np.log2(probabilities)).sum())
rng = np.random.default_rng(SEED)

names = ("Plug-in", "Miller--Madow", "Jackknife")
errors = {name: [] for name in names}
for sample_size in SAMPLE_SIZES:
    counts = rng.multinomial(sample_size, probabilities, size=TRIALS)
    for name, estimates in zip(names, entropy_estimates(counts), strict=True):
        errors[name].append(estimates - true_entropy)

colors = {
    "Plug-in": "#3F4650",
    "Miller--Madow": "#1E5A8A",
    "Jackknife": "#8B3A2B",
}
markers = {"Plug-in": "o", "Miller--Madow": "s", "Jackknife": "^"}

fig, ax = plt.subplots(figsize=(7.0, 4.2))
fig.patch.set_facecolor("#ffffff")
ax.set_facecolor("#ffffff")

for name in names:
    matrix = np.asarray(errors[name])
    means = matrix.mean(axis=1)
    lower, upper = np.percentile(matrix, [5, 95], axis=1)
    ax.fill_between(SAMPLE_SIZES, lower, upper, color=colors[name], alpha=0.09)
    ax.plot(
        SAMPLE_SIZES,
        means,
        marker=markers[name],
        color=colors[name],
        label=name,
        linewidth=1.8,
        markersize=5,
    )

ax.axhline(0, color="#30343B", linewidth=0.8)
ax.set_xscale("log", base=2)
ax.set_xticks(SAMPLE_SIZES)
ax.xaxis.set_major_formatter(ticker.FuncFormatter(lambda value, _: f"{int(value):,}"))
ax.set_xlabel("Sample size $n$", fontsize=12)
ax.set_ylabel("Entropy error (bits)", fontsize=12)
ax.legend(
    fontsize=10,
    frameon=True,
    fancybox=False,
    edgecolor="#8A8F98",
    loc="lower right",
)
ax.tick_params(labelsize=10, colors="#30343B")
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["bottom"].set_color("#3F4650")
ax.spines["left"].set_color("#3F4650")

fig.tight_layout()
validate_figure(fig)
output = figure_output_path(__file__, "convergence.png")
fig.savefig(output, dpi=300, bbox_inches="tight", facecolor="white")

print(
    f"seed={SEED} support={SUPPORT} s={ZIPF_EXPONENT} "
    f"trials={TRIALS} true_entropy_bits={true_entropy:.6f}"
)
print("| n | estimator | mean error (bits) | RMSE (bits) |")
print("|---:|:---|---:|---:|")
for row, sample_size in enumerate(SAMPLE_SIZES):
    for name in names:
        values = np.asarray(errors[name])[row]
        rmse = float(np.sqrt(np.mean(values**2)))
        print(f"| {sample_size} | {name} | {values.mean():+.4f} | {rmse:.4f} |")
print(f"saved={output}")