"""
04-04  Probability Simulator
Chapters 05-01, 05-02, 05-03

Run from the project folder:   python python/analysis.py
"""
from itertools import product
from math import comb, factorial, perm

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

rng = np.random.default_rng(2026)
rolls = pd.read_csv("data/dice_rolls.csv")
n = len(rolls)

# ── 1. EXACT DISTRIBUTION OF THE SUM OF TWO DICE ────────────────────────────
print("\n===== 1. THEORETICAL vs EMPIRICAL =====")

S = pd.DataFrame(list(product(range(1, 7), repeat=2)), columns=["d1", "d2"])
S["sum"] = S["d1"] + S["d2"]

ways = S["sum"].value_counts().sort_index()
theory = (ways / 36).to_numpy()
emp = (rolls["sum"].value_counts().reindex(range(2, 13), fill_value=0) / n).to_numpy()

comp = pd.DataFrame({
    "sum": range(2, 13),
    "ways": ways.to_numpy(),
    "theoretical": theory.round(4),
    "empirical": emp.round(4),
    "difference": (emp - theory).round(4),
})
print(comp.to_string(index=False))

print(f"\nCHECK sum of theoretical probabilities = {theory.sum():.4f}")
print(f"Largest gap between theory and {n} simulated rolls: {np.abs(emp - theory).max():.4f}")

# ── 2. PROBABILITY RULES ON THE DICE ────────────────────────────────────────
print("\n===== 2. PROBABILITY RULES =====")

pairs = [
    ("P(sum = 7)",        (S["sum"] == 7).mean(),  (rolls["sum"] == 7).mean()),
    ("P(sum >= 10)",      (S["sum"] >= 10).mean(), (rolls["sum"] >= 10).mean()),
    ("P(doubles)",        (S["d1"] == S["d2"]).mean(),
                          (rolls["die1"] == rolls["die2"]).mean()),
    ("P(at least one 6)", ((S["d1"] == 6) | (S["d2"] == 6)).mean(),
                          ((rolls["die1"] == 6) | (rolls["die2"] == 6)).mean()),
]
for label, t, e in pairs:
    print(f"{label:<20} theory {t:.4f}   empirical {e:.4f}")
print(f"  via the complement: 1 - (5/6)^2 = {1 - (5/6)**2:.4f}")

p7 = (S["sum"] == 7).mean()
pd_ = (S["d1"] == S["d2"]).mean()
pboth = ((S["sum"] == 7) & (S["d1"] == S["d2"])).mean()
print(f"P(sum = 7 OR doubles)  = {p7:.4f} + {pd_:.4f} - {pboth:.4f} = {p7 + pd_ - pboth:.4f}")
print("P(sum >= 10 | doubles) = "
      f"{((S['sum'] >= 10) & (S['d1'] == S['d2'])).mean() / pd_:.4f}")

# ── 3. LAW OF LARGE NUMBERS ─────────────────────────────────────────────────
print("\n===== 3. LAW OF LARGE NUMBERS =====")

running = np.cumsum(rolls["sum"] == 7) / np.arange(1, n + 1)
for k in (10, 50, 100, 500, 1000):
    print(f"  after {k:5d} rolls: {running[k-1]:.4f}   (error {abs(running[k-1] - 6/36):.4f})")

# ── 4. CONTINGENCY TABLE ────────────────────────────────────────────────────
print("\n===== 4. CONTINGENCY TABLE =====")

tab = pd.DataFrame([[45, 55, 20], [15, 30, 35]],
                   index=["Exercises", "Does not"],
                   columns=["Excellent", "Good", "Poor"])
with_margins = tab.copy()
with_margins["Total"] = with_margins.sum(axis=1)
with_margins.loc["Total"] = with_margins.sum(axis=0)
print(with_margins.to_string())

N = tab.values.sum()
print(f"\nmarginal    P(Exercises)              = {tab.loc['Exercises'].sum()/N:.4f}")
print(f"marginal    P(Excellent)              = {tab['Excellent'].sum()/N:.4f}")
print(f"joint       P(Exercises AND Excellent) = {tab.loc['Exercises','Excellent']/N:.4f}")
print(f"conditional P(Excellent | Exercises)  = "
      f"{tab.loc['Exercises','Excellent']/tab.loc['Exercises'].sum():.4f}  <- ROW total")
print(f"conditional P(Exercises | Excellent)  = "
      f"{tab.loc['Exercises','Excellent']/tab['Excellent'].sum():.4f}  <- COLUMN total")
union = (tab.loc["Exercises"].sum() + tab["Excellent"].sum()
         - tab.loc["Exercises", "Excellent"]) / N
print(f"union       P(Exercises OR Excellent) = {union:.4f}")

expected = np.outer(tab.sum(axis=1), tab.sum(axis=0)) / N
print("\nExpected counts IF independent:")
print(pd.DataFrame(expected, index=tab.index, columns=tab.columns).round(2).to_string())
print(f"Observed 45 vs expected {expected[0,0]:.1f}  ->  the variables are DEPENDENT")

# ── 5. BAYES' THEOREM ───────────────────────────────────────────────────────
print("\n===== 5. BAYES' THEOREM =====")

def bayes(prior, sens, fpr):
    evidence = sens * prior + fpr * (1 - prior)
    return evidence, sens * prior / evidence

sens, fpr = 0.95, 0.10
print(f"Test: {sens:.0%} sensitive, false-positive rate {fpr:.0%}\n")
print(" prevalence   P(positive)   P(disease | positive)")
for p in (0.001, 0.005, 0.01, 0.05, 0.10, 0.30, 0.50):
    ev, post = bayes(p, sens, fpr)
    print(f"   {p:6.3f}       {ev:6.4f}          {post:6.4f}")

p = 0.01
tp, fp = 10000 * p * sens, 10000 * (1 - p) * fpr
print(f"\nNatural frequencies at 1% prevalence, per 10,000 people:")
print(f"  with disease      {10000*p:6.0f}    of whom {tp:5.0f} test positive")
print(f"  without disease   {10000*(1-p):6.0f}    of whom {fp:5.0f} test positive")
print(f"  => P(disease | positive) = {tp:.0f} / {tp+fp:.0f} = {tp/(tp+fp):.4f}")

# ── 6. COUNTING ─────────────────────────────────────────────────────────────
print("\n===== 6. COUNTING =====")

print(f"5!                       = {factorial(5)}")
print(f"10P3 (order matters)     = {perm(10, 3)}")
print(f"10C3 (order does not)    = {comb(10, 3)}")
print(f"49C6 lottery combos      = {comb(49, 6):,}")
print(f"P(jackpot)               = {1/comb(49,6):.3e}")
p5 = comb(6, 5) * comb(43, 1) / comb(49, 6)
print(f"P(match exactly 5)       = {p5:.3e}  (1 in {round(1/p5):,})")
print(f"MISSISSIPPI arrangements = "
      f"{factorial(11)//(factorial(4)*factorial(4)*factorial(2)):,}")

def bday(k):
    return 1 - np.prod([(365 - i) / 365 for i in range(k)])

print("\nBirthday problem:")
for k in (10, 23, 30, 50, 70):
    print(f"  k = {k:2d} -> {bday(k):.4f}")

# ── 7. MONTY HALL ───────────────────────────────────────────────────────────
print("\n===== 7. MONTY HALL (100,000 simulated games) =====")

trials = 100_000
car = rng.integers(1, 4, trials)
pick = rng.integers(1, 4, trials)
print(f"  stay:   {(pick == car).mean():.4f}   (theory 1/3 = {1/3:.4f})")
print(f"  switch: {(pick != car).mean():.4f}   (theory 2/3 = {2/3:.4f})")
print("  Switching doubles your chance -- the host's choice is not independent")
print("  of where the car is, which is what makes the intuition fail.")

# ── 8. PLOTS ────────────────────────────────────────────────────────────────
fig, axes = plt.subplots(2, 2, figsize=(11, 8.5))

w = 0.4
xs = np.arange(2, 13)
axes[0, 0].bar(xs - w/2, theory, w, label="theoretical", color="#5B2A86")
axes[0, 0].bar(xs + w/2, emp, w, label="empirical", color="#0FA3A3")
axes[0, 0].set(title="Two dice: theory vs 1,000 rolls", xlabel="Sum", ylabel="Probability")
axes[0, 0].legend()

axes[0, 1].plot(running, color="#5B2A86")
axes[0, 1].axhline(6/36, color="#0FA3A3", lw=2, ls="--")
axes[0, 1].set(title="Law of Large Numbers", xlabel="Rolls",
               ylabel="Running P(sum = 7)", ylim=(0, 0.35))

prev = np.linspace(0.001, 0.5, 300)
post = [bayes(p, sens, fpr)[1] for p in prev]
axes[1, 0].plot(prev, post, color="#5B2A86", lw=2)
axes[1, 0].axvline(0.01, ls="--", color="grey")
axes[1, 0].set(title="Bayes: posterior vs base rate", xlabel="Prevalence",
               ylabel="P(disease | positive)")

ks = np.arange(1, 61)
axes[1, 1].plot(ks, [bday(k) for k in ks], color="#0B7A7A", lw=2)
axes[1, 1].axhline(0.5, ls="--", color="grey")
axes[1, 1].axvline(23, ls="--", color="grey")
axes[1, 1].set(title="Birthday problem", xlabel="People in the room",
               ylabel="P(at least one shared birthday)")

plt.tight_layout()
plt.savefig("probability_plots.png", dpi=110)
print("\nWrote probability_plots.png")