"""Reproduce the article's synthetic examples with Python 3 standard library only.

Run: python3 calculate.py
Writes calculation.json, model-ranges.csv and synthetic-log.csv beside this file.
No learner data, network calls, FSRS scheduler or random sampling are used.
"""

import csv
import json
import math
from decimal import Decimal, localcontext
from pathlib import Path


def distribution(n, p):
    """Calculate the exact binomial expression to 70-digit decimal precision."""
    with localcontext() as ctx:
        ctx.prec = 70
        return [
            Decimal(math.comb(n, k)) * p**k * (1 - p) ** (n - k) for k in range(n + 1)
        ]


def quantile(probs, threshold):
    """Return the smallest success count with CDF at least threshold."""
    total = Decimal(0)
    for k, probability in enumerate(probs):
        total += probability
        if total >= threshold:
            return k
    return len(probs) - 1


def summarize(n):
    p = Decimal("0.9")
    probs = distribution(n, p)
    lo, hi = [quantile(probs, Decimal(q)) for q in ("0.05", "0.95")]
    assert abs(sum(probs) - 1) < Decimal("1e-24")
    assert abs(sum(Decimal(k) * v for k, v in enumerate(probs)) - n * p) < Decimal(
        "1e-22"
    )
    assert abs(
        sum((Decimal(k) - n * p) ** 2 * v for k, v in enumerate(probs))
        - n * p * (1 - p)
    ) < Decimal("1e-20")
    return {
        "attempts": n,
        "assumed_p": float(p),
        "lower_count": lo,
        "upper_count": hi,
        "lower_percent": 100 * lo / n,
        "upper_percent": 100 * hi / n,
        "range_probability": float(sum(probs[lo : hi + 1])),
        "perfect_session_probability": float(probs[-1]),
        "sd_percentage_points": math.sqrt(float(p * (1 - p)) / n) * 100,
    }


rows = [summarize(n) for n in (20, 100, 1000)]
assert abs(rows[0]["perfect_session_probability"] - 0.9**20) < 1e-14
assert distribution(2, Decimal("0.5")) == [
    Decimal("0.25"),
    Decimal("0.50"),
    Decimal("0.25"),
]
# 同日の再回答を含む人工例。アプリから抽出した履歴ではありません。
log = [
    {
        "event": i,
        "card_id": f"DEMO-{i:02d}",
        "day": "illustrative-day-1",
        "attempt_for_card_today": 1,
        "outcome": "pass" if i <= 18 else "fail",
    }
    for i in range(1, 21)
]
log += [
    {
        "event": 20 + i,
        "card_id": f"DEMO-{19 + i % 2:02d}",
        "day": "illustrative-day-1",
        "attempt_for_card_today": 2 + (i - 1) // 2,
        "outcome": "pass",
    }
    for i in range(1, 21)
]
first = [x for x in log if x["attempt_for_card_today"] == 1]
assert len({x["card_id"] for x in log}) == 20
assert sum(x["outcome"] == "pass" for x in log) == 38
assert len(first) == 20 and sum(x["outcome"] == "pass" for x in first) == 18
out = Path(__file__).resolve().parent
result = {
    "label": "Synthetic arithmetic examples, not observed learner data or an FSRS forecast.",
    "assumptions": "Independent binary attempts with the same fixed p=0.9; counts are exact-model 5th and 95th percentile endpoints, not confidence intervals. Discreteness makes inclusive coverage exceed 90%.",
    "probability_method": "Binomial formula evaluated with Decimal precision 70; no RNG or seed.",
    "rows": rows,
    "denominator_example": {
        "distinct_cards": 20,
        "all_attempts": 40,
        "all_passes": 38,
        "all_attempt_pass_percent": 95,
        "first_attempts": 20,
        "first_passes": 18,
        "first_attempt_pass_percent": 90,
    },
    "checks": [
        "probability sums",
        "binomial mean and variance identities",
        "two fair trials control",
        "perfect-session identity",
        "synthetic-log accounting",
    ],
}
(out / "calculation.json").write_text(json.dumps(result, indent=2) + "\n")
for name, values in [("model-ranges.csv", rows), ("synthetic-log.csv", log)]:
    with (out / name).open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(values[0]), lineterminator="\n")
        writer.writeheader()
        writer.writerows(values)
print(json.dumps(result, indent=2))