"""Sample size calculator for parallel two-arm RCT (continuous outcome, e.g., ΔTIR).

Formula (two-sample, two-sided, equal allocation):
    N_per_arm = 2 * (z_{1-α/2} + z_{1-β})^2 * (SD/Δ)^2

Effect-size DB: Korean reference (placeholder values; replace with KDA/KAID data).
"""

from __future__ import annotations
import json
import math
import os

# Inverse normal CDF approximation (Beasley-Springer-Moro).
def _norm_inv(p: float) -> float:
    if p <= 0 or p >= 1:
        raise ValueError("p must be in (0,1)")
    a = [-3.969683028665376e+01, 2.209460984245205e+02, -2.759285104469687e+02,
         1.383577518672690e+02, -3.066479806614716e+01, 2.506628277459239e+00]
    b = [-5.447609879822406e+01, 1.615858368580409e+02, -1.556989798598866e+02,
         6.680131188771972e+01, -1.328068155288572e+01]
    c = [-7.784894002430293e-03, -3.223964580411365e-01, -2.400758277161838e+00,
         -2.549732539343734e+00, 4.374664141464968e+00, 2.938163982698783e+00]
    d = [7.784695709041462e-03, 3.224671290700398e-01, 2.445134137142996e+00,
         3.754408661907416e+00]
    p_low = 0.02425
    p_high = 1 - p_low
    if p < p_low:
        q = math.sqrt(-2 * math.log(p))
        return (((((c[0]*q + c[1])*q + c[2])*q + c[3])*q + c[4])*q + c[5]) / \
               ((((d[0]*q + d[1])*q + d[2])*q + d[3])*q + 1)
    if p <= p_high:
        q = p - 0.5
        r = q * q
        return (((((a[0]*r + a[1])*r + a[2])*r + a[3])*r + a[4])*r + a[5]) * q / \
               (((((b[0]*r + b[1])*r + b[2])*r + b[3])*r + b[4])*r + 1)
    q = math.sqrt(-2 * math.log(1 - p))
    return -(((((c[0]*q + c[1])*q + c[2])*q + c[3])*q + c[4])*q + c[5]) / \
           ((((d[0]*q + d[1])*q + d[2])*q + d[3])*q + 1)


def sample_size_two_arm(delta: float, sd: float, alpha: float = 0.05,
                        power: float = 0.80) -> dict:
    """Per-arm sample size for two-sample t-test on continuous outcome."""
    if delta <= 0 or sd <= 0:
        raise ValueError("delta and sd must be > 0")
    z_a = _norm_inv(1 - alpha / 2)
    z_b = _norm_inv(power)
    n = 2 * (z_a + z_b) ** 2 * (sd / delta) ** 2
    n_per_arm = math.ceil(n)
    return {
        "delta": delta,
        "sd": sd,
        "alpha": alpha,
        "power": power,
        "z_alpha_2": round(z_a, 4),
        "z_beta": round(z_b, 4),
        "n_per_arm": n_per_arm,
        "n_total": n_per_arm * 2,
        "formula": "N/arm = 2*(z_{1-a/2}+z_{1-b})^2 * (SD/delta)^2",
    }


def load_effect_size_db(path: str | None = None) -> dict:
    """Load Korean reference effect-size DB (JSON)."""
    if path is None:
        path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                            "data", "effect_sizes_kr.json")
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def by_hba1c_tertile(db: dict, alpha: float = 0.05, power: float = 0.80) -> list:
    """For each HbA1c tertile entry, compute N/arm using its effect/sd."""
    out = []
    for entry in db.get("hba1c_tertiles", []):
        try:
            res = sample_size_two_arm(entry["delta_TIR_pct"], entry["sd_TIR_pct"],
                                      alpha, power)
            out.append({**entry, **res})
        except Exception as e:
            out.append({**entry, "error": str(e)})
    return out