import pandas as pd, numpy as np
from scipy import stats
import statsmodels.api as sm
from coding import code, CAPABILITY, INSTITUTIONAL

pd.set_option("display.width", 200)

df = pd.read_csv("/mnt/user-data/uploads/1788949022214_draghi-measures.csv")
df["class"] = df["ID"].map(code)
df["impl"] = (df["Jul-26"] == "Implemented").astype(int)
df["impl_or_part"] = df["Jul-26"].isin(["Implemented", "Partially Implemented"]).astype(int)
order = {"Not Implemented": 0, "In Progress": 1, "Partially Implemented": 2, "Implemented": 3}
df["ord"] = df["Jul-26"].map(order)
df["ord0"] = df["Sep-25"].map(order)
df["moved"] = (df["ord"] > df["ord0"]).astype(int)

print("=" * 78)
print("1. CODING DISTRIBUTION")
print("=" * 78)
print(df["class"].value_counts().to_string())
print()
print(pd.crosstab(df["Sector"], df["class"]).to_string())

print()
print("=" * 78)
print("2. HEADLINE: implementation rate by class")
print("=" * 78)
tab = df.groupby("class").agg(
    n=("ID", "size"),
    implemented=("impl", "sum"),
    strict_rate=("impl", "mean"),
    impl_or_part=("impl_or_part", "sum"),
    full_partial_rate=("impl_or_part", "mean"),
    moved_since_sep25=("moved", "mean"),
)
print(tab.round(3).to_string())

# binary C vs K only
ck = df[df["class"].isin(["C", "K"])]
ct = pd.crosstab(ck["class"], ck["impl"])
print()
print("C vs K, strict implementation:")
print(ct.to_string())
odds, p = stats.fisher_exact(ct.values)
print(f"Fisher exact: OR(C over K) = {1/odds:.2f}  p = {p:.5f}")

ct2 = pd.crosstab(ck["class"], ck["impl_or_part"])
odds2, p2 = stats.fisher_exact(ct2.values)
print("C vs K, implemented-or-partial:")
print(ct2.to_string())
print(f"Fisher exact: OR(C over K) = {1/odds2:.2f}  p = {p2:.5f}")

print()
print("=" * 78)
print("3. THE SECTOR-SIZE / COMPOSITION CHECK")
print("=" * 78)
print("Strict implementation rate by sector AND class (n in parentheses):")
rows = []
for sec, g in df.groupby("Sector"):
    r = {"Sector": sec, "N": len(g)}
    for c in ["C", "K", "I"]:
        sub = g[g["class"] == c]
        r[c] = f"{sub['impl'].mean():.2f} ({len(sub)})" if len(sub) else "-"
    rows.append(r)
print(pd.DataFrame(rows).to_string(index=False))

print()
print("Within-sector paired comparison, C vs K (sectors with >=3 of each):")
diffs = []
for sec, g in df.groupby("Sector"):
    c = g[g["class"] == "C"]["impl"]
    k = g[g["class"] == "K"]["impl"]
    if len(c) >= 3 and len(k) >= 3:
        diffs.append((sec, len(c), c.mean(), len(k), k.mean(), c.mean() - k.mean()))
d = pd.DataFrame(diffs, columns=["Sector", "nC", "rateC", "nK", "rateK", "diff"])
print(d.round(3).to_string(index=False))
if len(d):
    print(f"\nMean within-sector difference: {d['diff'].mean():+.3f}")
    print(f"Sectors where C > K: {(d['diff'] > 0).sum()} of {len(d)}")
    try:
        w = stats.wilcoxon(d["diff"])
        print(f"Wilcoxon signed-rank on sector differences: p = {w.pvalue:.4f}")
    except Exception as e:
        print("Wilcoxon:", e)

print()
print("Cochran-Mantel-Haenszel (C vs K, stratified by sector):")
strata = []
for sec, g in df.groupby("Sector"):
    sub = g[g["class"].isin(["C", "K"])]
    if sub["class"].nunique() == 2:
        t = pd.crosstab(sub["class"], sub["impl"]).reindex(index=["C", "K"], columns=[0, 1]).fillna(0).values
        strata.append(t)
arr = np.stack(strata, axis=2)
cmh = sm.stats.StratifiedTable(arr)
print(f"  common OR (C over K) = {1/cmh.oddsratio_pooled:.3f}")
print(f"  CMH test of no association: p = {cmh.test_null_odds().pvalue:.5f}")
print(f"  test of OR homogeneity across sectors: p = {cmh.test_equal_odds().pvalue:.4f}")

print()
print("=" * 78)
print("4. LOGISTIC REGRESSION WITH SECTOR FIXED EFFECTS")
print("=" * 78)
X = pd.get_dummies(df[["class", "Sector"]], drop_first=True).astype(float)
X = sm.add_constant(X)
for dep, label in [("impl", "strict"), ("impl_or_part", "implemented-or-partial")]:
    print(f"\n--- outcome: {label} ---")
    m = sm.Logit(df[dep], X).fit(disp=0, method="bfgs", maxiter=500)
    keep = [i for i in m.params.index if i.startswith("class_")]
    for k in keep:
        print(f"  {k}: coef {m.params[k]:+.3f}  OR {np.exp(m.params[k]):.3f}  p {m.pvalues[k]:.4f}")
    # LR test: does class add anything beyond sector?
    Xs = sm.add_constant(pd.get_dummies(df[["Sector"]], drop_first=True).astype(float))
    m0 = sm.Logit(df[dep], Xs).fit(disp=0, method="bfgs", maxiter=500)
    lr = 2 * (m.llf - m0.llf)
    p = stats.chi2.sf(lr, len(keep))
    print(f"  LR test, class added to sector-only model: chi2={lr:.2f} df={len(keep)} p={p:.5f}")
    # reverse: does sector add anything beyond class?
    Xc = sm.add_constant(pd.get_dummies(df[["class"]], drop_first=True).astype(float))
    m1 = sm.Logit(df[dep], Xc).fit(disp=0, method="bfgs", maxiter=500)
    lr2 = 2 * (m.llf - m1.llf)
    p2 = stats.chi2.sf(lr2, 9)
    print(f"  LR test, sector added to class-only model: chi2={lr2:.2f} df=9 p={p2:.5f}")

print()
print("=" * 78)
print("5. IS THE SECTOR PATTERN A COMPOSITION EFFECT?")
print("=" * 78)
sec = df.groupby("Sector").agg(
    N=("ID", "size"),
    strict=("impl", "mean"),
    share_K=("class", lambda s: (s == "K").mean()),
    share_I=("class", lambda s: (s == "I").mean()),
    share_C=("class", lambda s: (s == "C").mean()),
).sort_values("strict", ascending=False)
print(sec.round(3).to_string())
print()
for v in ["share_K", "share_I", "share_C"]:
    r, p = stats.pearsonr(sec[v], sec["strict"])
    rs, ps = stats.spearmanr(sec[v], sec["strict"])
    print(f"  corr(sector {v}, sector strict rate): pearson r={r:+.3f} p={p:.4f} | spearman rho={rs:+.3f} p={ps:.4f}")
r, p = stats.pearsonr(sec["N"], sec["strict"])
print(f"  corr(sector size N, strict rate):        pearson r={r:+.3f} p={p:.4f}")
rs, ps = stats.spearmanr(sec["N"], sec["strict"])
print(f"                                           spearman rho={rs:+.3f} p={ps:.4f}")

print()
print("=" * 78)
print("6. TRAJECTORY BY CLASS (Sep-25 -> Jul-26)")
print("=" * 78)
for col in ["Sep-25", "Jan-26", "Jul-26"]:
    d2 = df.assign(x=(df[col] == "Implemented").astype(int),
                   y=df[col].isin(["Implemented", "Partially Implemented"]).astype(int))
    g = d2.groupby("class").agg(strict=("x", "mean"), full_partial=("y", "mean"))
    print(f"{col}: " + "  ".join(f"{c} strict={g.loc[c,'strict']:.3f} f+p={g.loc[c,'full_partial']:.3f}" for c in ["C", "K", "I"]))