"""Tolerance-matched runtime comparison: BBQ (EM) vs Crowd-BT (online gradient).

Both estimators are stopped by the SAME criterion -- the largest change in the
log-scale item score over one full sweep of the data -- so the seconds-vs-
minutes comparison of Figure 3 is measured rather than asserted.
"""
import json, math, time, sys
import numpy as np
import pandas as pd
from bbq_vs_crowdbt import load_clic, load_humaine, encode, ALPHA0, BETA0, MU0, SIGMA_SQ0, KAPPA

TOLS = [1e-2, 1e-3, 1e-4, 1e-5, 1e-6]


def bbq_timed(i_idx, j_idx, r_idx, K, R, a=5.0, b=0.1, alpha=10.0, beta=2.0, max_iter=4000):
    lam = np.ones(K); q = np.full(R, 0.5)
    w = np.ones(i_idx.size)
    n_r = np.bincount(r_idx, minlength=R).astype(float)
    hit = {}; t0 = time.perf_counter(); mind = np.inf; ll_prev = None
    for it in range(max_iter):
        y = lam[i_idx] / (lam[i_idx] + lam[j_idx])
        qc = q[r_idx]
        g = qc * y / (qc * y + (1.0 - qc) * 0.5)
        q_new = np.clip((np.bincount(r_idx, weights=g, minlength=R) + (alpha - 1.0))
                        / (n_r + alpha + beta - 2.0), 1e-12, 1 - 1e-12)
        wins = np.bincount(i_idx, weights=g, minlength=K) + (a - 1.0)
        den = np.zeros(K); s = lam[i_idx] + lam[j_idx]
        np.add.at(den, i_idx, g / s); np.add.at(den, j_idx, g / s)
        lam_new = wins / (den + b)
        shift = float(np.max(np.abs(np.log(lam_new) - np.log(lam))))
        yi = lam_new[i_idx] / (lam_new[i_idx] + lam_new[j_idx]); qq = q_new[r_idx]
        ll = float(np.sum(np.log(qq * yi + (1 - qq) * 0.5))
                   + np.sum((a - 1) * np.log(lam_new) - b * lam_new)
                   + np.sum((alpha - 1) * np.log(q_new) + (beta - 1) * np.log1p(-q_new)))
        if ll_prev is not None:
            mind = min(mind, ll - ll_prev)
        ll_prev = ll
        lam, q = lam_new, q_new
        for tol in TOLS:
            if tol not in hit and shift < tol:
                hit[tol] = (it + 1, time.perf_counter() - t0)
        if len(hit) == len(TOLS):
            break
    return lam, q, hit, it + 1, time.perf_counter() - t0, mind


def crowdbt_timed(i_idx, j_idx, r_idx, K, R, max_epochs=600, seed=0):
    mu = np.full(K, MU0); ss = np.full(K, SIGMA_SQ0)
    al = np.full(R, ALPHA0); be = np.full(R, BETA0)
    rng = np.random.default_rng(seed)
    order = np.arange(i_idx.size)
    hit = {}; t0 = time.perf_counter(); traj = []
    ii = i_idx.tolist(); jj = j_idx.tolist(); rr = r_idx.tolist()
    for ep in range(max_epochs):
        rng.shuffle(order)
        prev = mu.copy()
        for t in order:
            w = ii[t]; l = jj[t]; k = rr[t]
            mw = mu[w]; ml = mu[l]; sw = ss[w]; sl = ss[l]
            a = al[k]; b = be[k]
            m = mw - ml
            c1 = 1.0 if m > 30 else (0.0 if m < -30 else 1.0 / (1.0 + math.exp(-m)))
            c2 = 1.0 - c1
            ab = a + b
            dnm = a * c1 + b * c2
            if dnm <= 0:
                continue
            d = (a - b) * c1 * c2 / dnm
            mu[w] = mw + sw * d
            mu[l] = ml - sl * d
            h = (a - b) * c1 * c2 * (c2 - c1) / dnm - d * d
            fac = h - d * d
            ss[w] = sw * max(1.0 + sw * fac, KAPPA)
            ss[l] = sl * max(1.0 + sl * fac, KAPPA)
            f = a * c1 / dnm
            e1 = (f * (a + 1.0) + (1.0 - f) * a) / (ab + 1.0)
            e2 = (f * (a + 1.0) * (a + 2.0) + (1.0 - f) * a * (a + 1.0)) / ((ab + 1.0) * (ab + 2.0))
            v = e2 - e1 * e1
            if v > 1e-12 and e1 > e2 and 0.0 < e1 < 1.0:
                an = e1 * (e1 - e2) / v; bn = (1.0 - e1) * (e1 - e2) / v
                if an > 0 and bn > 0 and math.isfinite(an) and math.isfinite(bn):
                    al[k] = min(an, 1e6); be[k] = min(bn, 1e6)
        shift = float(np.max(np.abs(mu - prev)))
        traj.append(shift)
        for tol in TOLS:
            if tol not in hit and shift < tol:
                hit[tol] = (ep + 1, time.perf_counter() - t0)
        if len(hit) == len(TOLS):
            break
    return mu, al / (al + be), hit, ep + 1, time.perf_counter() - t0, traj


def run(name, df, max_epochs):
    i, j, r, K, R = encode(df)
    lam, q, bh, bi, bt, mind = bbq_timed(i, j, r, K, R)
    mu, eta, ch, ce, ct, traj = crowdbt_timed(i, j, r, K, R, max_epochs=max_epochs)
    from scipy.stats import kendalltau
    out = dict(dataset=name, comparisons=int(i.size), items=int(K), raters=int(R),
               bbq_total_iters=bi, bbq_total_secs=round(bt, 4),
               bbq_min_logpost_delta=mind,
               bbq_hit={str(k): [v[0], round(v[1], 4)] for k, v in bh.items()},
               cbt_total_epochs=ce, cbt_total_secs=round(ct, 3),
               cbt_hit={str(k): [v[0], round(v[1], 3)] for k, v in ch.items()},
               cbt_final_shift=round(traj[-1], 8),
               kendall_bbq_vs_cbt=round(float(kendalltau(np.log(lam), mu).statistic), 4))
    print(json.dumps(out), flush=True)
    return out


if __name__ == "__main__":
    which = sys.argv[1]
    out = []
    if which == "small":
        for s in ["screened", "unscreened"]:
            out.append(run("IHQ-" + s, load_clic(s), 4000))
        out.append(run("IHQ-all", pd.concat([load_clic("screened"), load_clic("unscreened")],
                                            ignore_index=True), 4000))
    elif which == "mtbench":
        from datasets import load_dataset
        ds = load_dataset("lmsys/mt_bench_human_judgments")
        rows = []
        for split in ds:
            for ex in ds[split]:
                w = ex.get("winner")
                if w == "model_a":
                    rows.append(dict(rater=str(ex.get("judge")), winner=ex["model_a"], loser=ex["model_b"]))
                elif w == "model_b":
                    rows.append(dict(rater=str(ex.get("judge")), winner=ex["model_b"], loser=ex["model_a"]))
        out.append(run("MT-Bench", pd.DataFrame(rows), 4000))
    else:
        h = load_humaine()
        sub = h.sample(n=105220, random_state=20260802).reset_index(drop=True)
        out.append(run("HUMAINE-105220", sub, 600))
    with open("timing2_%s.json" % which, "w") as f:
        json.dump(out, f, indent=1)