#!/usr/bin/env python3
"""MSC-P-011 two-arm binary-outcome sample-size reference. MIT; stdlib only."""

from __future__ import annotations

import csv
import math
from pathlib import Path
from statistics import NormalDist


DATA = Path(__file__).resolve().parents[1] / "datasets" / "msc-p011-ab-design.csv"


def main() -> None:
    rows = list(csv.DictReader(DATA.open(encoding="utf-8", newline="")))
    if len(rows) != 1:
        raise ValueError("reference dataset must contain one design")
    row = rows[0]
    p0, p1 = float(row["baseline_rate"]), float(row["target_rate"])
    alpha, power = float(row["alpha"]), float(row["power"])
    allocation = float(row["allocation_ratio"])
    attrition = float(row["attrition_rate"])
    design_effect = float(row["design_effect"])
    if not (0 < p0 < 1 and 0 < p1 < 1 and 0 < alpha < 1 and 0 < power < 1):
        raise ValueError("rates, alpha and power must be in (0,1)")
    if allocation != 1:
        raise ValueError("this reference formula supports equal allocation only")
    pooled = (p0 + p1) / 2
    z_alpha = NormalDist().inv_cdf(1 - alpha / 2)
    z_power = NormalDist().inv_cdf(power)
    raw = (
        z_alpha * math.sqrt(2 * pooled * (1 - pooled))
        + z_power * math.sqrt(p0 * (1 - p0) + p1 * (1 - p1))
    ) ** 2 / (p1 - p0) ** 2
    per_arm = math.ceil(raw)
    adjusted = math.ceil(raw * design_effect / (1 - attrition))
    total = 2 * adjusted
    print(f"design={row['design']}")
    print(f"absolute_mde={p1 - p0:.6f}")
    print(f"relative_lift={(p1 - p0) / p0:.6f}")
    print(f"z_alpha={z_alpha:.6f}")
    print(f"z_power={z_power:.6f}")
    print(f"raw_n_per_arm={raw:.6f}")
    print(f"rounded_n_per_arm={per_arm}")
    print(f"adjusted_n_per_arm={adjusted}")
    print(f"adjusted_total={total}")


if __name__ == "__main__":
    main()