# /// script
# requires-python = "==3.12.*"
# dependencies = ["matplotlib", "numpy"]
# ///
"""Compare Good--Turing estimated coverage with realized coverage.

The experiment uses repeated multinomial samples from a finite Zipf
distribution. Its finite support matters: both curves converge to one.
"""

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 = 42
SUPPORT = 5_000
ZIPF_EXPONENT = 1.1
TRIALS = 300
SAMPLE_SIZES = np.unique(np.round(np.geomspace(10, 10_000, 40)).astype(int))

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

estimated_mean = np.empty(len(SAMPLE_SIZES))
estimated_lo = np.empty(len(SAMPLE_SIZES))
estimated_hi = np.empty(len(SAMPLE_SIZES))
realized_mean = np.empty(len(SAMPLE_SIZES))

for i, sample_size in enumerate(SAMPLE_SIZES):
    counts = rng.multinomial(sample_size, probabilities, size=TRIALS)
    singleton_counts = np.count_nonzero(counts == 1, axis=1)
    estimated = 1.0 - singleton_counts / sample_size
    realized = (counts > 0) @ probabilities

    estimated_mean[i] = estimated.mean()
    estimated_lo[i], estimated_hi[i] = np.percentile(estimated, [5, 95])
    realized_mean[i] = realized.mean()

fig, ax = plt.subplots(figsize=(6.8, 4.0))
fig.patch.set_facecolor("#ffffff")
ax.set_facecolor("#ffffff")

ax.fill_between(
    SAMPLE_SIZES,
    estimated_lo,
    estimated_hi,
    color="#D9E7F2",
    linewidth=0,
    label=f"Good--Turing 90% interval ({TRIALS} trials)",
)
ax.plot(
    SAMPLE_SIZES,
    estimated_mean,
    color="#1E5A8A",
    linewidth=2.1,
    label="Mean Good--Turing estimate",
)
ax.plot(
    SAMPLE_SIZES,
    realized_mean,
    color="#8B3A2B",
    linewidth=1.8,
    linestyle="--",
    label="Mean realized coverage",
)
ax.axhline(
    1.0,
    color="#3F4650",
    linewidth=1.0,
    linestyle=":",
    label="Finite-support limit",
)

ax.set_xscale("log")
ax.set_xlim(9, 12_000)
ax.set_ylim(0, 1.04)
ax.set_xticks([10, 30, 100, 300, 1_000, 3_000, 10_000])
ax.xaxis.set_major_formatter(ticker.FuncFormatter(lambda value, _: f"{int(value):,}"))
ax.set_xlabel("Sample size $n$", fontsize=12)
ax.set_ylabel(r"Coverage", 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__, "completeness.png")
fig.savefig(output, dpi=300, bbox_inches="tight", facecolor="white")

reference_index = int(np.argmin(np.abs(SAMPLE_SIZES - 400)))
print(
    f"seed={SEED} support={SUPPORT} s={ZIPF_EXPONENT} trials={TRIALS}\n"
    f"n={SAMPLE_SIZES[reference_index]} "
    f"estimated_coverage={estimated_mean[reference_index]:.4f} "
    f"realized_coverage={realized_mean[reference_index]:.4f}\n"
    f"saved={output}"
)