#!/usr/bin/env python3
"""Recompute the baseline-choice examples with the Python standard library."""

from __future__ import annotations

import csv
import math
from pathlib import Path


SESSIONS = [85.0, 102.0, 98.0, 101.0, 99.0, 105.0]
WITHIN_SESSION_SD = 10.0
THRESHOLDS = [10.0, 15.0]
BASELINE_SIZES = [1, 3, 5]


def normal_upper_tail(value: float, sd: float) -> float:
    return 0.5 * math.erfc(value / (sd * math.sqrt(2.0)))


def rows() -> list[dict[str, float | int]]:
    future = SESSIONS[-1]
    output: list[dict[str, float | int]] = []
    for n in BASELINE_SIZES:
        baseline = sum(SESSIONS[:n]) / n
        change_sd = WITHIN_SESSION_SD * math.sqrt(1.0 + 1.0 / n)
        output.append(
            {
                "baseline_sessions": n,
                "worked_baseline": baseline,
                "future_session": future,
                "worked_apparent_change": future - baseline,
                "model_change_sd": change_sd,
                "p_gain_at_least_10": normal_upper_tail(THRESHOLDS[0], change_sd),
                "p_gain_at_least_15": normal_upper_tail(THRESHOLDS[1], change_sd),
            }
        )
    return output


def main() -> None:
    computed = rows()
    expected = {
        1: (85.0, 20.0, 14.142136, 0.239750, 0.144422),
        3: (95.0, 10.0, 11.547005, 0.193238, 0.096965),
        5: (97.0, 8.0, 10.954451, 0.180655, 0.085452),
    }
    for row in computed:
        n = int(row["baseline_sessions"])
        actual = (
            float(row["worked_baseline"]),
            float(row["worked_apparent_change"]),
            float(row["model_change_sd"]),
            float(row["p_gain_at_least_10"]),
            float(row["p_gain_at_least_15"]),
        )
        for value, target in zip(actual, expected[n], strict=True):
            assert math.isclose(value, target, abs_tol=5e-7), (n, value, target)

    destination = Path(__file__).with_name("examples.csv")
    fieldnames = list(computed[0])
    with destination.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in computed:
            writer.writerow(
                {
                    **row,
                    "worked_baseline": f'{float(row["worked_baseline"]):.1f}',
                    "future_session": f'{float(row["future_session"]):.1f}',
                    "worked_apparent_change": f'{float(row["worked_apparent_change"]):.1f}',
                    "model_change_sd": f'{float(row["model_change_sd"]):.4f}',
                    "p_gain_at_least_10": f'{float(row["p_gain_at_least_10"]):.6f}',
                    "p_gain_at_least_15": f'{float(row["p_gain_at_least_15"]):.6f}',
                }
            )
    print(f"verified {len(computed)} baseline rules; wrote {destination.name}")


if __name__ == "__main__":
    main()