"""
04-08  Confidence Interval Estimator
Chapters 09-01, 09-02

Run from the project folder:   python python/analysis.py
"""
import numpy as np
import pandas as pd
from scipy import stats
from statsmodels.stats.proportion import proportion_confint
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

rng = np.random.default_rng(2026)
battery = pd.read_csv("data/battery_life.csv")["hours"].to_numpy()
poll = pd.read_csv("data/poll_responses.csv")["supports"]


# ── REUSABLE ESTIMATORS ─────────────────────────────────────────────────────
def ci_mean_t(x, conf=0.95):
    x = np.asarray(x)
    n = x.size
    se = x.std(ddof=1) / np.sqrt(n)
    tc = stats.t.ppf(1 - (1 - conf) / 2, n - 1)
    return pd.Series({"n": n, "mean": x.mean(), "sd": x.std(ddof=1), "se": se,
                      "df": n - 1, "t_crit": tc, "margin": tc * se,
                      "lower": x.mean() - tc * se, "upper": x.mean() + tc * se,
                      "width": 2 * tc * se})


def ci_mean_z(x, sigma, conf=0.95):
    x = np.asarray(x)
    n = x.size
    se = sigma / np.sqrt(n)
    zc = stats.norm.ppf(1 - (1 - conf) / 2)
    return pd.Series({"n": n, "mean": x.mean(), "sigma": sigma, "se": se,
                      "z_crit": zc, "margin": zc * se,
                      "lower": x.mean() - zc * se, "upper": x.mean() + zc * se})


def ci_prop(x, n, conf=0.95):
    ph = x / n
    se = np.sqrt(ph * (1 - ph) / n)
    zc = stats.norm.ppf(1 - (1 - conf) / 2)
    return pd.Series({"x": x, "n": n, "p_hat": ph, "se": se, "z_crit": zc,
                      "margin": zc * se, "lower": ph - zc * se,
                      "upper": ph + zc * se, "width": 2 * zc * se})


# ── 1. CI FOR A MEAN ────────────────────────────────────────────────────────
print("\n===== 1. CONFIDENCE INTERVAL FOR A MEAN — battery_life.csv =====")
print("sigma is UNKNOWN (s was computed from the data)  ->  use t\n")
print(ci_mean_t(battery).round(4).to_string())

print("\nscipy agrees:")
print(stats.t.interval(0.95, battery.size - 1, battery.mean(), stats.sem(battery)))

print("\nIf sigma were known to be 96 (z-interval, for contrast):")
print(ci_mean_z(battery, sigma=96).round(4).to_string())

# ── 2. THREE CONFIDENCE LEVELS ──────────────────────────────────────────────
print("\n===== 2. THE CONFIDENCE / PRECISION TRADE-OFF =====")

levels = [0.90, 0.95, 0.99]
tab = pd.DataFrame([ci_mean_t(battery, cl)[["t_crit", "margin", "lower",
                                            "upper", "width"]]
                    for cl in levels], index=[f"{int(cl*100)}%" for cl in levels])
print(tab.round(4).to_string())
print("\nHigher confidence -> larger critical value -> WIDER interval.")
print("You cannot have high confidence AND a narrow interval AND a small sample.")

# ── 3. CI FOR A PROPORTION ──────────────────────────────────────────────────
print("\n===== 3. CONFIDENCE INTERVAL FOR A PROPORTION — poll_responses.csv =====")

x = int((poll == "Yes").sum())
n = int(poll.size)
print(f"condition check: n*p_hat = {x}  and  n*(1-p_hat) = {n-x}   (both >= 5)\n")
print(ci_prop(x, n).round(5).to_string())

print("\nAlternative methods on the SAME data:")
for name, method in [("Wald", "normal"), ("Wilson score", "wilson"),
                     ("Clopper-Pearson", "beta"), ("Agresti-Coull", "agresti_coull")]:
    lo, hi = proportion_confint(x, n, method=method)
    print(f"  {name:<16} {lo:.4f} to {hi:.4f}")
print("\nAt n = 600 all four agree closely. At small n they diverge sharply.")

# ── 4. SAMPLE SIZE ──────────────────────────────────────────────────────────
print("\n===== 4. SAMPLE SIZE FOR A TARGET MARGIN OF ERROR =====")


def n_mean(sigma, E, conf=0.95):
    return int(np.ceil((stats.norm.ppf(1 - (1 - conf) / 2) * sigma / E) ** 2))


def n_prop(E, p=0.5, conf=0.95):
    return int(np.ceil(p * (1 - p) * (stats.norm.ppf(1 - (1 - conf) / 2) / E) ** 2))


s_pilot = battery.std(ddof=1)
print(f"Using the pilot sd = {s_pilot:.2f} hours\n")
print("  MEAN                 90%      95%      99%")
for E in (50, 25, 10, 5):
    print(f"  E = {E:3d} hours     {n_mean(s_pilot,E,0.90):6d}   "
          f"{n_mean(s_pilot,E,0.95):6d}   {n_mean(s_pilot,E,0.99):6d}")

ph = x / n
print(f"\n  PROPORTION        p=0.5    p={ph:.2f}")
for E in (0.05, 0.03, 0.02, 0.01):
    print(f"  E = {E:.2f}          {n_prop(E):6d}   {n_prop(E, ph):6d}")
print("\nHalving E QUADRUPLES n -- precision is bought with the square of the sample.")

# ── 5. CI FOR THE VARIANCE ──────────────────────────────────────────────────
print("\n===== 5. CONFIDENCE INTERVAL FOR sigma^2 AND sigma =====")

nb = battery.size
s2 = battery.var(ddof=1)
df = nb - 1
lo = df * s2 / stats.chi2.ppf(0.975, df)
hi = df * s2 / stats.chi2.ppf(0.025, df)
print(f"s^2 = {s2:.3f}   df = {df}")
print(f"95% CI for sigma^2:  {lo:.3f} to {hi:.3f}")
print(f"95% CI for sigma  :  {np.sqrt(lo):.3f} to {np.sqrt(hi):.3f}")
print(f"\nDistance below s^2: {s2-lo:.2f}    above: {hi-s2:.2f}   -> NOT symmetric,")
print("because the chi-square distribution is right-skewed.")
print("This interval also REQUIRES normality and is not robust to it.")

# ── 6. COVERAGE SIMULATION ──────────────────────────────────────────────────
print("\n===== 6. DOES A 95% INTERVAL REALLY COVER 95% OF THE TIME? =====")

true_mu, true_sigma, R = 1218, 96, 5000

hits = 0
for _ in range(R):
    smp = rng.normal(true_mu, true_sigma, 45)
    r = ci_mean_t(smp)
    hits += r["lower"] <= true_mu <= r["upper"]
print(f"t-interval for a mean, n = 45:  coverage = {hits/R:.4f}  (target 0.95)")


def cover_p(n, p, R=5000):
    ph_ = rng.binomial(n, p, R) / n
    e = stats.norm.ppf(0.975) * np.sqrt(ph_ * (1 - ph_) / n)
    return np.mean((ph_ - e <= p) & (p <= ph_ + e))


print("\nWald proportion interval coverage, by n and p:")
rows = [(nn, pp, round(cover_p(nn, pp), 4))
        for nn in (20, 100, 600) for pp in (0.05, 0.50)]
print(pd.DataFrame(rows, columns=["n", "p", "coverage"]).to_string(index=False))
print("\nWald delivers its promise near p = 0.5, but UNDER-covers badly near 0 or 1.")
print("That is why prop.test and statsmodels default to the WILSON interval.")

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

K = 100
los, his = [], []
for _ in range(K):
    smp = rng.normal(true_mu, true_sigma, 45)
    r = ci_mean_t(smp)
    los.append(r["lower"]); his.append(r["upper"])
los, his = np.array(los), np.array(his)
hit = (los <= true_mu) & (true_mu <= his)
for i in range(K):
    axes[0, 0].plot([los[i], his[i]], [i, i],
                    color="#0FA3A3" if hit[i] else "#b4122e", lw=2)
axes[0, 0].axvline(true_mu, color="#5B2A86", lw=2)
axes[0, 0].set(title=f"100 x 95% CIs — {(~hit).sum()} missed",
               xlabel="hours", ylabel="simulated study")

axes[0, 1].bar([f"{int(cl*100)}%" for cl in levels], tab["width"], color="#5B2A86")
axes[0, 1].set(title="Interval width vs confidence", ylabel="width (hours)")

Es = np.arange(5, 61)
axes[1, 0].plot(Es, [n_mean(s_pilot, E) for E in Es], color="#0B7A7A", lw=2)
axes[1, 0].set_yscale("log")
axes[1, 0].set(title="Sample size vs margin of error",
               xlabel="target E (hours)", ylabel="required n (log scale)")

r = ci_mean_t(battery)
axes[1, 1].hist(battery, bins=12, color="#8A5FBF", edgecolor="white")
axes[1, 1].axvline(r["mean"], color="#5B2A86", lw=2)
axes[1, 1].axvline(r["lower"], color="#0FA3A3", lw=2, ls="--")
axes[1, 1].axvline(r["upper"], color="#0FA3A3", lw=2, ls="--")
axes[1, 1].set(title="Battery life with 95% CI for the mean", xlabel="hours")

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