"""Estimate booking pressure under explicit, uncalibrated renewal assumptions."""
import argparse
import json
import math
from pathlib import Path

import numpy as np

from pipeline import read_config


def simulate(data):
    assumptions = data["planning_assumptions"]
    rng = np.random.default_rng(assumptions["seed"])
    n = assumptions["incoming_learners"]
    p_reschedule = assumptions["reschedule_probability"]
    lessons = data["lessons_per_package"]
    results = []
    for p_renewal in assumptions["renewal_probabilities"]:
        renewals = rng.binomial(n, p_renewal, assumptions["trials"])
        original_slots = (n + renewals) * lessons
        total = original_slots + rng.binomial(original_slots, p_reschedule)
        expected = n * (1 + p_renewal) * lessons * (1 + p_reschedule)
        mean = float(total.mean())
        assert abs(mean - expected) < 5, "Simulation mean differs from analytical expectation"
        p50, p90 = (int(value) for value in np.quantile(total, [0.50, 0.90], method="higher"))
        results.append({
            "renewal_probability": p_renewal,
            "simulated_mean": round(mean, 2), "analytical_mean": round(expected, 2),
            "median_slots": p50, "p90_slots": p90,
            "weekly_slots_for_p90": math.ceil(p90 / assumptions["teaching_weeks"]),
            "overflow": [
                {"slots_per_week": capacity, "annual_slots": capacity * assumptions["teaching_weeks"],
                 "probability": round(float(np.mean(total > capacity * assumptions["teaching_weeks"])), 4)}
                for capacity in assumptions["weekly_slot_options"]
            ],
        })
    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("inputs", type=Path)
    parser.add_argument("output", type=Path)
    options = parser.parse_args()
    options.output.write_text(json.dumps(simulate(read_config(options.inputs)), indent=2) + "\n", encoding="utf-8")