#!/usr/bin/env python3
"""Reproduce the personal-best selection simulation.

Python standard library only. Writes results.csv beside this script unless an
output path is supplied. No participant or product data are read.
"""

from __future__ import annotations

import csv
import math
import random
import statistics
import sys
from pathlib import Path

ATTEMPTS = (1, 2, 3, 5, 10, 20, 50)
REPLICATES = 250_000
MODELS = (
    ("normal", 20260930),
    ("bounded_uniform", 20260931),
    ("student_t_5", 20260932),
)


def draw(kind: str, rng: random.Random) -> float:
    if kind == "normal":
        return rng.gauss(0.0, 1.0)
    if kind == "bounded_uniform":
        return rng.uniform(-math.sqrt(3.0), math.sqrt(3.0))
    if kind == "student_t_5":
        normal = rng.gauss(0.0, 1.0)
        chi_square = rng.gammavariate(2.5, 2.0)
        # A t(5) variable has variance 5/3; rescale to variance 1.
        return (normal / math.sqrt(chi_square / 5.0)) * math.sqrt(3.0 / 5.0)
    raise ValueError(kind)


def quantile(values: list[float], probability: float) -> float:
    ordered = sorted(values)
    position = (len(ordered) - 1) * probability
    low = int(position)
    high = min(low + 1, len(ordered) - 1)
    weight = position - low
    return ordered[low] * (1.0 - weight) + ordered[high] * weight


def simulate(kind: str, seed: int) -> list[dict[str, str | int]]:
    rng = random.Random(seed)
    bests = {n: [] for n in ATTEMPTS}
    means = {n: [] for n in ATTEMPTS}
    for _ in range(REPLICATES):
        running_best = -math.inf
        running_sum = 0.0
        for attempt in range(1, max(ATTEMPTS) + 1):
            score = draw(kind, rng)
            running_best = max(running_best, score)
            running_sum += score
            if attempt in bests:
                bests[attempt].append(running_best)
                means[attempt].append(running_sum / attempt)

    rows: list[dict[str, str | int]] = []
    for attempts in ATTEMPTS:
        records = bests[attempts]
        rows.append(
            {
                "model": kind,
                "seed": seed,
                "replicates": REPLICATES,
                "attempts": attempts,
                "mean_personal_best_z": f"{statistics.mean(records):.6f}",
                "median_personal_best_z": f"{statistics.median(records):.6f}",
                "p2_5_personal_best_z": f"{quantile(records, 0.025):.6f}",
                "p97_5_personal_best_z": f"{quantile(records, 0.975):.6f}",
                "mean_all_attempts_z": f"{statistics.mean(means[attempts]):.6f}",
                "mc_se_mean_best": f"{statistics.stdev(records) / math.sqrt(REPLICATES):.6f}",
            }
        )
    return rows


def main() -> None:
    output = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).with_name("results.csv")
    rows = [row for kind, seed in MODELS for row in simulate(kind, seed)]
    with output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    print(f"wrote {len(rows)} rows to {output}")


if __name__ == "__main__":
    main()