"""
04-07  Sampling & CLT Simulator
Chapters 08-01, 08-02

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

rng = np.random.default_rng(2026)
pop = pd.read_csv("data/population.csv")
N = len(pop)
v = pop["value"].to_numpy()

# ── 1. TRUE POPULATION PARAMETERS ───────────────────────────────────────────
print("\n===== 1. POPULATION PARAMETERS (known exactly) =====")

mu = v.mean()
sigma = v.std(ddof=0)          # POPULATION sd: divide by N, not n-1

print(f"N            {N}")
print(f"mu           {mu:.4f}")
print(f"sigma        {sigma:.4f}    (population formula, ddof=0)")
print(f"sample sd    {v.std(ddof=1):.4f}    (n-1 version, for contrast)")
print(f"skewness     {stats.skew(v):+.4f}   -> the population is right-skewed")

print("\nBy region:")
print(pop.groupby("region")["value"].agg(["count", "mean", "std"]).round(3).to_string())

# ── 2. FOUR DESIGNS, ONE SAMPLE EACH (n = 40) ───────────────────────────────
print("\n===== 2. FOUR SAMPLING DESIGNS, n = 40 =====")

n = 40

srs = pop.sample(n, random_state=1)

k = N // n
start = int(rng.integers(0, k))
sys = pop.iloc[start::k].head(n)

alloc = (n * pop["region"].value_counts() / N).round().astype(int)
strat = pd.concat([g.sample(alloc[r], random_state=2)
                   for r, g in pop.groupby("region")])

pop["cluster"] = np.repeat(np.arange(1, 51), 20)      # 50 clusters of 20
chosen = rng.choice(np.arange(1, 51), 2, replace=False)
clus = pop[pop["cluster"].isin(chosen)]

print(f"proportional allocation: "
      f"{', '.join(f'{r}={c}' for r, c in alloc.items())}  (sums to {alloc.sum()})")

one_shot = pd.DataFrame({
    "design": ["SRS", "Systematic", "Stratified", "Cluster"],
    "n": [len(srs), len(sys), len(strat), len(clus)],
    "mean": [srs["value"].mean(), sys["value"].mean(),
             strat["value"].mean(), clus["value"].mean()],
})
one_shot["error"] = (one_shot["mean"] - mu).abs()
print(one_shot.round(3).to_string(index=False))
print("\nOne sample proves nothing -- repeat each design 2,000 times.")

# ── 3. PRECISION COMPARISON (2,000 replications) ────────────────────────────
print("\n===== 3. PRECISION COMPARISON (2,000 replications each) =====")

R = 2000

srs_means = rng.choice(v, size=(R, n)).mean(axis=1)

sys_means = np.array([v[int(rng.integers(0, k))::k][:n].mean() for _ in range(R)])

idx_by_region = {r: g.index.to_numpy() for r, g in pop.groupby("region")}
strat_means = np.array([
    v[np.concatenate([rng.choice(idx_by_region[r], alloc[r], replace=False)
                      for r in alloc.index])].mean()
    for _ in range(R)])

cluster_ids = pop["cluster"].to_numpy()
clus_means = np.array([
    v[np.isin(cluster_ids, rng.choice(np.arange(1, 51), 2, replace=False))].mean()
    for _ in range(R)])

prec = pd.DataFrame({
    "design": ["SRS", "Systematic", "Stratified", "Cluster"],
    "mean_of_means": [srs_means.mean(), sys_means.mean(),
                      strat_means.mean(), clus_means.mean()],
    "bias": [srs_means.mean() - mu, sys_means.mean() - mu,
             strat_means.mean() - mu, clus_means.mean() - mu],
    "std_error": [srs_means.std(ddof=1), sys_means.std(ddof=1),
                  strat_means.std(ddof=1), clus_means.std(ddof=1)],
})
prec["relative_to_SRS"] = prec["std_error"] / prec["std_error"].iloc[0]
print(prec.round(4).to_string(index=False))

print(f"\ntheoretical SRS standard error  sigma/sqrt(n) = {sigma/np.sqrt(n):.4f}")
print("Every design is UNBIASED (bias ~ 0) but their PRECISION differs.")
print("Stratified wins because the regional means genuinely differ; cluster")
print("loses because each draw commits to whole regions at a time.")

# ── 4. CENTRAL LIMIT THEOREM ────────────────────────────────────────────────
print("\n===== 4. CENTRAL LIMIT THEOREM =====")

sizes = [1, 5, 30, 100]
clt = {m: rng.choice(v, size=(3000, m)).mean(axis=1) for m in sizes}

clt_tab = pd.DataFrame({
    "n": sizes,
    "sim_mean": [clt[m].mean() for m in sizes],
    "sim_se": [clt[m].std(ddof=1) for m in sizes],
    "theory_se": [sigma / np.sqrt(m) for m in sizes],
    "sim_skew": [stats.skew(clt[m]) for m in sizes],
})
print(clt_tab.round(4).to_string(index=False))

print("\nThree things to notice:")
print(" 1. sim_mean is mu at EVERY n -- the sample mean is unbiased.")
print(" 2. sim_se tracks sigma/sqrt(n) almost exactly.")
print(" 3. sim_skew falls toward 0 as n grows -- that IS the CLT.")

# ── 5. BIAS: A CONVENIENCE SAMPLE ───────────────────────────────────────────
print("\n===== 5. WHY A BIGGER SAMPLE CANNOT FIX BIAS =====")

srt = np.sort(v)[::-1]
conv40, conv400 = srt[:40].mean(), srt[:400].mean()
print(f"true mu                        {mu:.3f}")
print(f"convenience sample of 40       {conv40:.3f}   (bias {conv40-mu:+.3f})")
print(f"convenience sample of 400      {conv400:.3f}   (bias {conv400-mu:+.3f})")
print("The bigger convenience sample is still wrong -- just more precisely wrong.")

# ── 6. PLOTS ────────────────────────────────────────────────────────────────
fig, axes = plt.subplots(2, 3, figsize=(14, 8.5))
ax = axes.ravel()

ax[0].hist(v, bins=40, color="#8A5FBF")
ax[0].axvline(mu, color="#0FA3A3", lw=2)
ax[0].set(title="Population (right-skewed)", xlabel="value")

for i, m in enumerate(sizes, start=1):
    ax[i].hist(clt[m], bins=40, color="#0FA3A3", range=(clt[1].min(), clt[1].max()))
    ax[i].axvline(mu, color="#5B2A86", lw=2)
    ax[i].set(title=f"Sampling distribution, n = {m}", xlabel="sample mean")

ax[5].boxplot([srs_means, sys_means, strat_means, clus_means],
              tick_labels=["SRS", "Systematic", "Stratified", "Cluster"])
ax[5].axhline(mu, ls="--", color="grey")
ax[5].set(title="Precision by design (2,000 reps)", ylabel="sample mean")
ax[5].tick_params(axis="x", rotation=30)

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