from pathlib import Path
import numpy as np
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf

ROOT = Path("analysis/designed-difficulty")
df = pd.read_csv(ROOT / "baseline_joined.csv")

df = df[
    (df["variant"] == "multiturn")
    & (df["turn"] == 0)
    & (df["status"] == "ok")
    & (df["dropped"] == 0)
    & df["choice"].isin(["hold", "diverge", "abstain"])
].copy()

df["y"] = (df["choice"] == "diverge").astype(int)

df["framing"] = df["framing"].fillna("unspecified")
df["category"] = df["category"].fillna("unspecified")

# Repeated epochs, shared questions across models, and
# questions sharing a claim are clustered together.
df["cluster"] = df["claim_id"].fillna(
    df["question_id"]
)

print("Observations:", len(df))
print("Questions:", df["question_id"].nunique())
print("Claims:", df["cluster"].nunique())
print("Overall divergence:", round(df["y"].mean(), 4))

formulas = {
    "A": (
        "y ~ evidence_scarcity + question_openness "
        "+ C(model)"
    ),
    "B": (
        "y ~ evidence_scarcity + question_openness "
        "+ C(model) + C(category) + C(framing)"
    ),
}

results = []

for name, formula in formulas.items():
    print(f"\n=== MODEL {name} ===")

    model = smf.glm(
        formula,
        data=df,
        family=sm.families.Binomial()
    )

    fit = model.fit(
        cov_type="cluster",
        cov_kwds={"groups": df["cluster"]}
    )

    for variable in [
        "evidence_scarcity",
        "question_openness",
    ]:
        coef = fit.params[variable]
        se = fit.bse[variable]
        pval = fit.pvalues[variable]

        odds_ratio = np.exp(coef)
        ci_low = np.exp(coef - 1.96 * se)
        ci_high = np.exp(coef + 1.96 * se)

        results.append({
            "model": name,
            "predictor": variable,
            "coefficient": coef,
            "odds_ratio": odds_ratio,
            "ci_low": ci_low,
            "ci_high": ci_high,
            "p_value": pval,
        })

        print(
            f"{variable:20s} "
            f"OR={odds_ratio:.3f} "
            f"95% CI=[{ci_low:.3f}, {ci_high:.3f}] "
            f"p={pval:.4g}"
        )

    print("AIC:", round(fit.aic, 2))
    print("Converged:", fit.converged)

output = ROOT / "regression_results.csv"
pd.DataFrame(results).to_csv(output, index=False)

print("\nSaved:", output)