import pandas as pd
import numpy as np
from scipy import stats
from scipy.stats import chi2_contingency
import warnings
warnings.filterwarnings('ignore')

df = pd.read_csv('06_Data_Deidentified_reviewer.csv')
results = []

def check(label, computed, expected, tol=0.01):
    if isinstance(expected, str):
        match = str(computed) == expected
    elif isinstance(expected, float):
        match = abs(float(computed) - expected) < tol
    else:
        match = computed == expected
    status = "PASS" if match else "FAIL"
    results.append((status, label, computed, expected))
    return match

# ============ 1. SAMPLE CHARACTERISTICS (Table 1) ============
N = len(df)
check("N total", N, 16075)

male = (df['gender'] == 0).sum()
female = (df['gender'] == 1).sum()
check("Male n", male, 6727)
check("Male %", round(male/N*100, 1), 41.8)
check("Female n", female, 9348)
check("Female %", round(female/N*100, 1), 58.2)

age_valid = df['age'].dropna()
check("Age mean", round(age_valid.mean(), 1), 20.1)
check("Age SD", round(age_valid.std(), 2), 2.69)
check("Age missing", df['age'].isnull().sum(), 284)

urban = (df['origin'] == 'urban').sum()
rural = (df['origin'] == 'rural').sum()
check("Urban n", urban, 5670)
check("Urban %", round(urban/N*100, 1), 35.3)
check("Rural n", rural, 10405)

dfc = (df['tier_code'] == 1).sum()
prov_dfc = (df['tier_code'] == 2).sum()
regular = (df['tier_code'] == 3).sum()
voc = (df['tier_code'] == 4).sum()
check("DFC n", dfc, 424)
check("Prov DFC n", prov_dfc, 983)
check("Regular n", regular, 14392)
check("Vocational n", voc, 276)

elite_n = df['elite'].sum()
non_elite_n = N - elite_n
check("Elite n", elite_n, 1407)
check("Elite %", round(elite_n/N*100, 1), 8.8)
check("Non-elite n", non_elite_n, 14668)

climate_missing = df['climate'].isnull().sum()
check("Climate missing", climate_missing, 1791)
check("Climate missing %", round(climate_missing/N*100, 1), 11.1)

path_df = df[df['climate'].notna()]
path_n = len(path_df)
check("Path sample N", path_n, 14284)
path_elite = path_df['elite'].sum()
check("Path elite n", path_elite, 1275)
check("Path elite %", round(path_elite/path_n*100, 1), 8.9)

all_items = [f'item_{i}' for i in range(1, 26)]
full_endorsing_mask = (df[all_items] == 5).all(axis=1)
full_endorsing_n = full_endorsing_mask.sum()
check("Full-endorsing n", full_endorsing_n, 1022)
check("Full-endorsing %", round(full_endorsing_n/N*100, 1), 6.4)

fe_in_path = (full_endorsing_mask & df['climate'].notna()).sum()
exclusion_n = path_n - fe_in_path
check("Exclusion sample N", exclusion_n, 13374)
check("FE in path sample", fe_in_path, 910)
check("FE not in path (missing climate)", full_endorsing_n - fe_in_path, 112)

print("=== Section 1: Sample Characteristics (Table 1) ===")
for s, l, c, e in results:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

results2 = []

# ============ 2. VR/IR STATISTICS (Section 3.2) ============
VR_items = [1, 3, 5, 6, 7, 12, 14, 15, 17, 19, 22, 23, 24, 25]
IR_items = [2, 4, 8, 9, 10, 11, 13, 16, 18, 20, 21]
VR_cols = [f'item_{i}' for i in VR_items]
IR_cols = [f'item_{i}' for i in IR_items]

check2 = lambda l, c, e, t=0.01: results2.append(("PASS" if abs(c-e)<t else "FAIL", l, c, e))

check2("VR mean", round(df['VR_mean'].mean(), 2), 4.35)
check2("VR SD", round(df['VR_mean'].std(), 2), 0.63)
check2("IR mean", round(df['IR_mean'].mean(), 2), 3.33)
check2("IR SD", round(df['IR_mean'].std(), 2), 0.70)

def cronbach_alpha(data):
    k = data.shape[1]
    var_sum = data.var(axis=0, ddof=1).sum()
    total_var = data.sum(axis=1).var(ddof=1)
    return (k / (k - 1)) * (1 - var_sum / total_var)

vr_alpha = cronbach_alpha(df[VR_cols])
ir_alpha = cronbach_alpha(df[IR_cols])
check2("VR alpha", round(vr_alpha, 3), 0.955)
check2("IR alpha", round(ir_alpha, 3), 0.865)

r_vi = df['VR_mean'].corr(df['IR_mean'])
check2("VR-IR r", round(r_vi, 2), 0.26)

check2("Satisfaction mean", round(df['satisfaction'].mean(), 2), 3.54)
check2("Satisfaction SD", round(df['satisfaction'].std(), 2), 0.89)
check2("Climate mean", round(df['climate'].mean(), 2), 3.52)
check2("Climate SD", round(df['climate'].std(), 2), 0.85)

print("\n=== Section 2: VR/IR Statistics (Section 3.2) ===")
for s, l, c, e in results2:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

# ============ 3. LPA PROFILE STATISTICS (Table 3) ============
results3 = []
check3 = lambda l, c, e, t=0.01: results3.append(("PASS" if abs(c-e)<t else "FAIL", l, c, e))

expected_profiles = [
    (1, 1022, 6.4, 5.00, 5.00, 4.11, 4.11, 1.09, 17.1),
    (2, 1864, 11.6, 5.00, 3.39, 3.85, 3.79, 1.31, 8.6),
    (3, 5196, 32.3, 4.74, 3.12, 3.59, 3.54, 1.22, 7.3),
    (4, 4416, 27.5, 4.15, 3.24, 3.43, 3.38, 1.14, 7.5),
    (5, 1689, 10.5, 4.00, 3.53, 3.45, 3.43, 1.11, 7.0),
    (6, 1888, 11.7, 3.08, 2.99, 3.10, 3.26, 0.98, 12.7),
]

for p in range(1, 7):
    pdata = df[df['profile'] == p]
    n = len(pdata)
    vr = pdata['VR_mean'].mean()
    ir = pdata['IR_mean'].mean()
    sat = pdata['satisfaction'].mean()
    clim = pdata['climate'].mean()
    cour = pdata['courses'].mean()
    elite_pct = pdata['elite'].mean() * 100
    exp = expected_profiles[p-1]
    check3(f"P{p} n", n, exp[1])
    check3(f"P{p} %", round(n/N*100, 1), exp[2])
    check3(f"P{p} VR", round(vr, 2), exp[3])
    check3(f"P{p} IR", round(ir, 2), exp[4])
    check3(f"P{p} Sat", round(sat, 2), exp[5])
    check3(f"P{p} Climate", round(clim, 2), exp[6])
    check3(f"P{p} Courses", round(cour, 2), exp[7])
    check3(f"P{p} Elite%", round(elite_pct, 1), exp[8])

print("\n=== Section 3: LPA Profiles (Table 3) ===")
for s, l, c, e in results3:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

# ============ 4. INSTITUTION COMPARISONS (Table 5) ============
results4 = []
check4 = lambda l, c, e, t=0.01: results4.append(("PASS" if abs(c-e)<t else "FAIL", l, c, e))

elite_df = df[df['elite'] == 1]
ordinary_df = df[df['elite'] == 0]

# Cohen's d helper: sample-size weighted pooled SD (standard formula,
# matching the manuscript Table 5 and Section 4.2 of the reproducibility doc)
def cohens_d(g1, g2):
    n1, n2 = len(g1), len(g2)
    s1, s2 = g1.std(ddof=1), g2.std(ddof=1)
    pooled_sd = np.sqrt(((n1-1)*s1**2 + (n2-1)*s2**2) / (n1+n2-2))
    return (g1.mean() - g2.mean()) / pooled_sd

# Satisfaction
elite_sat = elite_df['satisfaction'].mean()
ordinary_sat = ordinary_df['satisfaction'].mean()
t_stat, _ = stats.ttest_ind(elite_df['satisfaction'].dropna(), ordinary_df['satisfaction'].dropna(), equal_var=False)
d = cohens_d(elite_df['satisfaction'].dropna(), ordinary_df['satisfaction'].dropna())
check4("Sat elite M", round(elite_sat, 3), 3.641)
check4("Sat ordinary M", round(ordinary_sat, 3), 3.528)
check4("Sat t", round(t_stat, 2), 3.88)
check4("Sat d", round(d, 2), 0.13)

# Climate
elite_clim = path_df[path_df['elite'] == 1]['climate'].mean()
ordinary_clim = path_df[path_df['elite'] == 0]['climate'].mean()
t_stat2, _ = stats.ttest_ind(path_df[path_df['elite']==1]['climate'].dropna(), path_df[path_df['elite']==0]['climate'].dropna(), equal_var=False)
d2 = cohens_d(path_df[path_df['elite']==1]['climate'].dropna(), path_df[path_df['elite']==0]['climate'].dropna())
check4("Climate elite M", round(elite_clim, 3), 3.723)
check4("Climate ordinary M", round(ordinary_clim, 3), 3.503)
check4("Climate t", round(t_stat2, 2), 7.85)
check4("Climate d", round(d2, 2), 0.26)

# VR
elite_vr = elite_df['VR_mean'].mean()
ordinary_vr = ordinary_df['VR_mean'].mean()
t_stat3, p_val3 = stats.ttest_ind(elite_df['VR_mean'], ordinary_df['VR_mean'], equal_var=False)
d3 = cohens_d(elite_df['VR_mean'], ordinary_df['VR_mean'])
check4("VR elite M", round(elite_vr, 3), 4.290)
check4("VR ordinary M", round(ordinary_vr, 3), 4.357)
check4("VR t", round(t_stat3, 2), -3.23)
check4("VR p", round(p_val3, 3), 0.001)
check4("VR d", round(d3, 2), -0.11)

# IR
elite_ir = elite_df['IR_mean'].mean()
ordinary_ir = ordinary_df['IR_mean'].mean()
t_stat4, _ = stats.ttest_ind(elite_df['IR_mean'], ordinary_df['IR_mean'], equal_var=False)
d4 = cohens_d(elite_df['IR_mean'], ordinary_df['IR_mean'])
check4("IR elite M", round(elite_ir, 3), 3.467)
check4("IR ordinary M", round(ordinary_ir, 3), 3.318)
check4("IR t", round(t_stat4, 2), 6.53)
check4("IR d", round(d4, 2), 0.21)

print("\n=== Section 4: Institution Comparisons (Table 5) ===")
for s, l, c, e in results4:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

# ============ 5. FULL-ENDORISING COMPOSITION ============
results5 = []
check5 = lambda l, c, e, t=0.01: results5.append(("PASS" if abs(c-e)<t else "FAIL", l, c, e))

vr5_mask = df['VR_mean'] >= 4.999999
vr5_n = vr5_mask.sum()
check5("VR=5 n", vr5_n, 2886)
check5("VR=5 %", round(vr5_n/N*100, 1), 18.0)

fe_among_vr5 = (full_endorsing_mask & vr5_mask).sum()
check5("FE among VR=5", fe_among_vr5, 1022)
check5("FE among VR=5 %", round(fe_among_vr5/vr5_n*100, 1), 35.4)

ve_among_vr5 = vr5_n - fe_among_vr5
check5("VE among VR=5", ve_among_vr5, 1864)
check5("VE among VR=5 %", round(ve_among_vr5/vr5_n*100, 1), 64.6)

elite_vr5_n = ((df['elite']==1) & vr5_mask).sum()
check5("Elite VR=5 n", elite_vr5_n, 336)
elite_vr5_fe = ((df['elite']==1) & vr5_mask & full_endorsing_mask).sum()
check5("Elite VR=5 FE %", round(elite_vr5_fe/elite_vr5_n*100, 1), 52.1)

nonelite_vr5_n = ((df['elite']==0) & vr5_mask).sum()
check5("Non-elite VR=5 n", nonelite_vr5_n, 2550)
nonelite_vr5_fe = ((df['elite']==0) & vr5_mask & full_endorsing_mask).sum()
check5("Non-elite VR=5 FE %", round(nonelite_vr5_fe/nonelite_vr5_n*100, 1), 33.2)

contingency = [[elite_vr5_fe, elite_vr5_n - elite_vr5_fe],
               [nonelite_vr5_fe, nonelite_vr5_n - nonelite_vr5_fe]]
chi2_val, _, _, _ = chi2_contingency(contingency)
check5("Chi2 FE composition", round(chi2_val, 1), 45.4)

ir5_mask = df['IR_mean'] >= 4.999999
ir5_n = ir5_mask.sum()
check5("IR=5 n", ir5_n, 1057)
check5("IR=5 %", round(ir5_n/N*100, 1), 6.6)
ir5_fe = (ir5_mask & full_endorsing_mask).sum()
check5("IR=5 FE %", round(ir5_fe/ir5_n*100, 1), 96.7)

print("\n=== Section 5: Full-Endorsing Composition ===")
for s, l, c, e in results5:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

# ============ 6. DOSE-RESPONSE (Table 8) ============
results6 = []
check6 = lambda l, c, e, t=0.01: results6.append(("PASS" if abs(c-e)<t else "FAIL", l, c, e))

expected_data = [(4526, 4.290, 3.394), (6800, 4.346, 3.296), (2391, 4.371, 3.306), (2358, 4.465, 3.335)]
for c_val in [0, 1, 2, 3]:
    cdata = df[df['courses'] == c_val]
    cn = len(cdata)
    cvr = cdata['VR_mean'].mean()
    cir = cdata['IR_mean'].mean()
    check6(f"Courses={c_val} n", cn, expected_data[c_val][0])
    check6(f"Courses={c_val} VR", round(cvr, 3), expected_data[c_val][1])
    check6(f"Courses={c_val} IR", round(cir, 3), expected_data[c_val][2])

r_vr_c = df['VR_mean'].corr(df['courses'])
r_ir_c = df['IR_mean'].corr(df['courses'])
check6("r(VR, courses)", round(r_vr_c, 3), 0.086)
check6("r(IR, courses)", round(r_ir_c, 3), -0.029)

df_excl = df[~full_endorsing_mask]
r_vr_c_excl = df_excl['VR_mean'].corr(df_excl['courses'])
check6("r(VR, courses) after exclusion", round(r_vr_c_excl, 3), 0.098)

print("\n=== Section 6: Dose-Response (Table 8) ===")
for s, l, c, e in results6:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

# ============ 7. FOUR-QUADRANT (Table 9) ============
results7 = []
check7 = lambda l, c, e, t=0.1: results7.append(("PASS" if abs(c-e)<t else "FAIL", l, c, e))

vr_median = df['VR_mean'].median()
ir_median = df['IR_mean'].median()
check7("VR median", round(vr_median, 2), 4.50)
check7("IR median", round(ir_median, 2), 3.18)

vr_high = df['VR_mean'] > vr_median
ir_high = df['IR_mean'] > ir_median

q1 = (vr_high & ir_high).sum()
q2 = (vr_high & ~ir_high).sum()
q3 = (~vr_high & ir_high).sum()
q4 = (~vr_high & ~ir_high).sum()

check7("Q1 high-high %", round(q1/N*100, 1), 23.6)
check7("Q2 value-dominant %", round(q2/N*100, 1), 23.8)
check7("Q3 inst-dominant %", round(q3/N*100, 1), 23.8)
check7("Q4 low-low %", round(q4/N*100, 1), 28.8)

quad_labels = np.select([vr_high & ir_high, vr_high & ~ir_high, ~vr_high & ir_high, ~vr_high & ~ir_high], [0, 1, 2, 3])
quad_elite_ct = pd.crosstab(pd.Categorical(quad_labels), df['elite'])
chi2_quad, _, _, _ = chi2_contingency(quad_elite_ct)
check7("Quadrant chi2", round(chi2_quad, 1), 28.2)

expected_shares = [10.6, 7.4, 9.1, 8.0]
for qi, qmask in enumerate([vr_high & ir_high, vr_high & ~ir_high, ~vr_high & ir_high, ~vr_high & ~ir_high]):
    elite_in_q = df[qmask]['elite'].sum()
    elite_share = elite_in_q / qmask.sum() * 100
    check7(f"Q{qi+1} elite share %", round(elite_share, 1), expected_shares[qi])

print("\n=== Section 7: Four-Quadrant (Table 9) ===")
for s, l, c, e in results7:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

# ============ 8. RESPONSE-STYLE TIERS (Table 9b) ============
results8 = []
check8 = lambda l, c, e, t=0.1: results8.append(("PASS" if abs(c-e)<t else "FAIL", l, c, e))

max_items = (df[all_items] == 5).sum(axis=1)
t1_mask = max_items == 25
t2_mask = (max_items >= 1) & (max_items < 25)
t3_mask = max_items == 0

check8("T1 n", t1_mask.sum(), 1022)
check8("T1 %", round(t1_mask.sum()/N*100, 1), 6.4)
check8("T2 n", t2_mask.sum(), 11766)
check8("T2 %", round(t2_mask.sum()/N*100, 1), 73.2)
check8("T3 n", t3_mask.sum(), 3287)
check8("T3 %", round(t3_mask.sum()/N*100, 1), 20.4)

t1_elite = df[t1_mask]['elite'].mean() * 100
check8("T1 elite %", round(t1_elite, 1), 17.1)
t2_elite = df[t2_mask]['elite'].mean() * 100
check8("T2 elite %", round(t2_elite, 1), 7.8)
t3_elite = df[t3_mask]['elite'].mean() * 100
check8("T3 elite %", round(t3_elite, 1), 9.5)

tier_labels = np.select([t1_mask, t2_mask, t3_mask], [0, 1, 2])
tier_elite_ct = pd.crosstab(pd.Categorical(tier_labels), df['elite'])
chi2_tier, _, _, _ = chi2_contingency(tier_elite_ct)
check8("Tier chi2", round(chi2_tier, 1), 105.2)

print("\n=== Section 8: Response-Style Tiers (Table 9b) ===")
for s, l, c, e in results8:
    print(f"  [{s}] {l}: computed={c}, expected={e}")

# ============ SUMMARY ============
all_results = results + results2 + results3 + results4 + results5 + results6 + results7 + results8
total = len(all_results)
passed = sum(1 for r in all_results if r[0] == "PASS")
failed = sum(1 for r in all_results if r[0] == "FAIL")
print(f"\n{'='*60}")
print(f"SUMMARY: {passed}/{total} PASS, {failed} FAIL")
if failed > 0:
    print("\nFAILED checks:")
    for s, l, c, e in all_results:
        if s == "FAIL":
            print(f"  [FAIL] {l}: computed={c}, expected={e}")
print(f"{'='*60}")