"""Claim 6: BBQ rater-quality parameter q_r vs rater agreement (IHQ unscreened).

Measures the Pearson correlation between the EM-fitted mixture weight q_r
(Eq. 2/Eq. 12) and each rater's agreement with the final ranking, under nine
definitions of 'agreement' and three choices of 'final ranking'.  Adds a
destructive control (permuted q_r) and the operational test the claim makes:
whether q_r alone separates the hand-screened raters from the unscreened ones
on the pooled IHQ-all data.
"""
import json
import numpy as np
from scipy import stats
from bbq_vs_crowdbt import bbq_em, load_clic, encode


def consensus(df):
    i, j, r, K, R = encode(df)
    lam, q, its, hist = bbq_em(np.ones(i.size), i, j, r, K, R)
    items = sorted(set(df.winner) | set(df.loser))
    return items, lam, q, {v: k for k, v in enumerate(items)}, (i, j, r)


def rater_agreement(df, rank_of, mode):
    """Per-rater agreement with a reference ranking (higher rank value = better)."""
    out = {}
    for rater, g in df.groupby("rater"):
        w = g.winner.to_numpy()
        l = g.loser.to_numpy()
        rw = np.array([rank_of[x] for x in w])
        rl = np.array([rank_of[x] for x in l])
        keep = rw != rl
        if mode == "all_pairs":
            correct = (rw > rl)
            n = correct.size
        else:                                # drop reference-ties
            correct = (rw[keep] > rl[keep])
            n = correct.size
        out[rater] = np.nan if n == 0 else correct.mean()
    return out


def main():
    res = {}
    scr = load_clic("screened")
    uns = load_clic("unscreened")
    allc = __import__("pandas").concat([scr, uns], ignore_index=True)

    # --- fits
    it_u, lam_u, q_u, map_u, _ = consensus(uns)
    it_s, lam_s, q_s, map_s, _ = consensus(scr)
    it_a, lam_a, q_a, map_a, _ = consensus(allc)

    refs = {
        "BBQ consensus on IHQ-unscreened": (it_u, lam_u),
        "BBQ consensus on IHQ-screened":   (it_s, lam_s),
        "BBQ consensus on IHQ-all":        (it_a, lam_a),
    }
    raters_u = sorted(set(uns.rater))
    qmap_u = dict(zip(sorted(set(uns.rater)), q_u))
    rows = []
    for rname, (items, lam) in refs.items():
        rank_of = {it: float(l) for it, l in zip(items, lam)}
        for mode in ["all_pairs", "drop_ref_ties"]:
            agree = rater_agreement(uns, rank_of, mode)
            xs = np.array([qmap_u[r] for r in raters_u])
            ys = np.array([agree[r] for r in raters_u])
            m = np.isfinite(ys)
            pr = stats.pearsonr(xs[m], ys[m])
            sp = stats.spearmanr(xs[m], ys[m])
            rows.append(dict(reference=rname, agreement=mode, n_raters=int(m.sum()),
                             pearson_r=round(float(pr.statistic), 4),
                             pearson_p=float(pr.pvalue),
                             spearman_rho=round(float(sp.statistic), 4)))
    res["correlations"] = rows

    # --- same, on the screened split (paper says unscreened shows the effect)
    qmap_s = dict(zip(sorted(set(scr.rater)), q_s))
    rank_of = {it: float(l) for it, l in zip(it_s, lam_s)}
    ag_s = rater_agreement(scr, rank_of, "all_pairs")
    rs = sorted(set(scr.rater))
    xs = np.array([qmap_s[r] for r in rs]); ys = np.array([ag_s[r] for r in rs])
    m = np.isfinite(ys)
    res["screened_split_pearson_r"] = round(float(stats.pearsonr(xs[m], ys[m]).statistic), 4)
    res["screened_split_n"] = int(m.sum())

    # --- destructive control: permute q_r
    rank_of = {it: float(l) for it, l in zip(it_u, lam_u)}
    agree = rater_agreement(uns, rank_of, "all_pairs")
    xs = np.array([qmap_u[r] for r in raters_u])
    ys = np.array([agree[r] for r in raters_u])
    m = np.isfinite(ys)
    rng = np.random.default_rng(20260802)
    perm = [float(stats.pearsonr(rng.permutation(xs[m]), ys[m]).statistic) for _ in range(2000)]
    res["permutation_control"] = dict(
        observed=round(float(stats.pearsonr(xs[m], ys[m]).statistic), 4),
        perm_mean=round(float(np.mean(perm)), 4),
        perm_abs_max=round(float(np.max(np.abs(perm))), 4),
        perm_p=float((np.sum(np.array(perm) >= stats.pearsonr(xs[m], ys[m]).statistic) + 1) / 2001))

    # --- operational test: does q_r alone recover the hand-screening label?
    scr_set, uns_set = set(scr.rater), set(uns.rater)
    raters_a = sorted(set(allc.rater))
    qmap_a = dict(zip(raters_a, q_a))
    lab = np.array([1 if r in scr_set else 0 for r in raters_a])
    qv = np.array([qmap_a[r] for r in raters_a])
    u = stats.mannwhitneyu(qv[lab == 1], qv[lab == 0], alternative="greater")
    auc = float(u.statistic / ((lab == 1).sum() * (lab == 0).sum()))
    res["screening_recovery"] = dict(
        n_screened=int((lab == 1).sum()), n_unscreened=int((lab == 0).sum()),
        mean_q_screened=round(float(qv[lab == 1].mean()), 4),
        mean_q_unscreened=round(float(qv[lab == 0].mean()), 4),
        auc_q_predicts_screened=round(auc, 4), mwu_p=float(u.pvalue))

    # --- operational test 2: drop the lowest-q raters from unscreened and see
    # whether the resulting consensus moves toward the hand-screened consensus.
    base_rank = np.argsort(np.argsort(-np.array([float(l) for l in lam_s])))
    s_order = {it: float(l) for it, l in zip(it_s, lam_s)}
    def tau_to_screened(sub):
        items, lam, q, _, _ = consensus(sub)
        common = [x for x in items if x in s_order]
        a = np.array([float(lam[items.index(x)]) for x in common])
        b = np.array([s_order[x] for x in common])
        return float(stats.kendalltau(a, b).statistic)
    drops = []
    order = np.argsort(np.array([qmap_u[r] for r in raters_u]))
    for k in [0, 5, 10, 15, 20]:
        bad = set(np.array(raters_u)[order[:k]])
        sub = uns[~uns.rater.isin(bad)]
        drops.append(dict(dropped=k, kept_raters=int(sub.rater.nunique()),
                          comparisons=int(len(sub)),
                          kendall_to_screened=round(tau_to_screened(sub), 4)))
    res["low_quality_drop"] = drops
    res["claimed_r"] = 0.724
    print(json.dumps(res, indent=1))
    with open("rater_quality.json", "w") as f:
        json.dump(res, f, indent=1)


if __name__ == "__main__":
    main()