import numpy as np
import pandas as pd
from pathlib import Path

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

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

df["y"] = df.choice.eq("diverge").astype(int)
df["group"] = np.where(df.category.eq("other"), 1, 0)
df["cluster"] = df.claim_id.fillna(df.question_id)

models = sorted(df.model.unique())
claims = sorted(df.cluster.unique())

arr_n = np.zeros((len(claims), len(models), 2, 2))
arr_y = np.zeros_like(arr_n)

ci = {c: i for i, c in enumerate(claims)}
mi = {m: i for i, m in enumerate(models)}

for r in df.itertuples():
    idx = (
        ci[r.cluster],
        mi[r.model],
        int(r.group),
        int(r.evidence_scarcity)
    )
    arr_n[idx] += 1
    arr_y[idx] += r.y

def compute(weights):
    n = np.einsum("i,imjk->mjk", weights, arr_n)
    y = np.einsum("i,imjk->mjk", weights, arr_y)
    if (n == 0).any():
        return None
    rates = y / n
    rd = rates[:, :, 1] - rates[:, :, 0]
    return (rd[:, 1] - rd[:, 0]) * 100

observed = compute(np.ones(len(claims)))

rng = np.random.default_rng(42)
boot = []

for _ in range(4000):
    weights = rng.multinomial(
        len(claims),
        np.ones(len(claims)) / len(claims)
    )
    value = compute(weights)
    if value is not None:
        boot.append(value)

boot = np.array(boot)

rows = []
for i, model in enumerate(models):
    low, high = np.percentile(boot[:, i], [2.5, 97.5])
    rows.append({
        "model": model,
        "group_difference_pp": observed[i],
        "ci_low": low,
        "ci_high": high,
        "bootstrap_n": len(boot)
    })

result = pd.DataFrame(rows)
result.to_csv(out / "bootstrap_s01_results.csv", index=False)
print(result.to_string(index=False))