"""Shared loading and statistics for the Role-Play Bench audit (pure Python, no third-party packages)."""
import csv, datetime, json, math, random, statistics as st
from collections import defaultdict, Counter
A = '${RPB_AUDIT_DIR}'
RAW, OUT = f'{A}/raw', f'{A}/out'
DIMS = ['basics', 'logic', 'knowledge', 'diversity', 'content_logic', 'interaction']
# Role-Play Bench weights (worker/src/benchmarks/minimax-rp/index.ts groups/dimensions; scoring.ts groupScores)
GROUPS = [('worlds', 0.5, {'basics': 1, 'logic': 1, 'knowledge': 1}), ('stories', 0.25, {'diversity': 4, 'content_logic': 2}), ('preferences', 0.25, {'interaction': 1})]
# Published leaderboard (index.ts referenceLeaderboard, dataset card 2026-01)
PUBLISHED = {'MiniMax-M2-her': 84.65, 'doubao-1.5-pro': 80.64, 'gpt-5.1': 80.63, 'claude-opus-4.5-highthinking': 76.62, 'gemini-3-pro': 75.6,
             'claude-opus-4.5-lowthinking': 71.19, 'claude-sonnet-4.5': 69.35, 'gemini-2.5-pro': 68.23, 'deepseek-v3.1': 64.22, 'deepseek-v3.2': 60.27, 'grok-4.1': 48.47}
MM_RUN = {k: ('mm-' + k) if k != 'MiniMax-M2-her' else 'mm-minimax-m2-her' for k in PUBLISHED}

def utc(ms):
    return datetime.datetime.fromtimestamp(ms / 1000, datetime.timezone.utc).strftime('%Y-%m-%d %H:%M:%S') if ms else ''

def ov(d):
    """Overall from dimension scores, as scoring.ts groupScores (groups with no finite member are skipped)."""
    tot = ws = 0.0
    for _, w, members in GROUPS:
        m = {k: v for k, v in members.items() if d.get(k) is not None and math.isfinite(d[k])}
        if not m: continue
        tot += w * sum(v * d[k] for k, v in m.items()) / sum(m.values()); ws += w
    return tot / ws if ws else float('nan')

def agg(rows):
    """scoring.ts summarize: mean of each dimension over sessions, then the group/overall formula."""
    dims = {}
    for k in DIMS:
        v = [r[k] for r in rows if r.get(k) is not None and math.isfinite(r[k])]
        if v: dims[k] = sum(v) / len(v)
    return dims, ov(dims)

def load_jsonl(name):
    return [json.loads(l) for l in open(f'{RAW}/export/{name}.jsonl')]

def load_runs():
    meta = {x['id']: x for x in json.load(open(f'{RAW}/runs_meta.json'))[0]['results']}
    for x in json.load(open(f'{RAW}/runs_detail.json'))[0]['results']:
        m = meta[x['id']]
        m['config'] = json.loads(x['config_json']); m['judges'] = json.loads(x['judges']) if x['judges'] else []
        m['usersim'] = json.loads(x['usersim']) if x['usersim'] else {}; m['pv'] = x['pv']
        m['result'] = json.loads(x['result_json']) if x['result_json'] else None; m['dispute_cfg'] = json.loads(x['dispute']) if x['dispute'] else None
        m['engine'] = json.loads(x['engine']) if x['engine'] else None
    return meta

def panel(run):
    js = run['judges']; keys = sorted((j['key'], j['model'], ((j.get('extra') or {}).get('reasoning') or {}).get('effort')) for j in js)
    if run['origin'] == 'minimax': return 'imported (MiniMax judge, unpublished)'
    if keys == [('judge_a', 'anthropic/claude-opus-5.5', 'high'), ('judge_b', 'openai/gpt-6.1-sol', 'medium')]: return 'official'
    if {k for k, _, _ in keys} <= {'judge_w', 'judge_s'}: return 'screening'
    return 'other'

def judges_label(run):
    return ' + '.join(f"{j['key']}={j['model'].split('/')[-1]}" + (f"({((j.get('extra') or {}).get('reasoning') or {}).get('effort')})" if ((j.get('extra') or {}).get('reasoning') or {}).get('effort') else '') + f"×{j.get('weight')}" for j in run['judges'])

# ---------- statistics ----------
def mean(x): return sum(x) / len(x) if x else float('nan')

def pearson(x, y):
    if len(x) < 3: return float('nan')
    mx, my = mean(x), mean(y); sx = sum((a - mx) ** 2 for a in x); sy = sum((b - my) ** 2 for b in y)
    return sum((a - mx) * (b - my) for a, b in zip(x, y)) / math.sqrt(sx * sy) if sx and sy else float('nan')

def ranks(x):
    o = sorted(range(len(x)), key=lambda i: x[i]); r = [0.0] * len(x); i = 0
    while i < len(o):
        j = i
        while j + 1 < len(o) and x[o[j + 1]] == x[o[i]]: j += 1
        for k in range(i, j + 1): r[o[k]] = (i + j) / 2 + 1
        i = j + 1
    return r

def spearman(x, y): return pearson(ranks(x), ranks(y))

def icc21(x, y):
    """ICC(2,1): two-way random effects, absolute agreement, single rater (Shrout & Fleiss)."""
    n, k = len(x), 2
    if n < 3: return float('nan')
    rows = list(zip(x, y)); gm = mean(x + y)
    ssr = k * sum((mean(r) - gm) ** 2 for r in rows); ssc = n * ((mean(x) - gm) ** 2 + (mean(y) - gm) ** 2)
    sst = sum((v - gm) ** 2 for r in rows for v in r); sse = sst - ssr - ssc
    msr, msc, mse = ssr / (n - 1), ssc / (k - 1), sse / ((n - 1) * (k - 1))
    den = msr + (k - 1) * mse + k * (msc - mse) / n
    return (msr - mse) / den if den else float('nan')

def kappa(x, y):
    n = len(x); po = sum(a == b for a, b in zip(x, y)) / n
    cats = set(x) | set(y); pe = sum((x.count(c) / n) * (y.count(c) / n) for c in cats)
    return (po - pe) / (1 - pe) if pe < 1 else float('nan'), po

def kendall_tau_b(x, y):
    n = len(x); c = d = tx = ty = 0
    for i in range(n):
        for j in range(i + 1, n):
            a, b = x[i] - x[j], y[i] - y[j]
            if a == 0 and b == 0: continue
            if a == 0: tx += 1
            elif b == 0: ty += 1
            elif a * b > 0: c += 1
            else: d += 1
    den = math.sqrt((c + d + tx) * (c + d + ty))
    return (c - d) / den if den else float('nan')

def kendall_w(matrix):
    """matrix: list of raters (scenarios), each a list of scores for the same m items; higher = better. Tie-corrected W."""
    m, n = len(matrix), len(matrix[0])
    R = [0.0] * n; T = 0.0
    for row in matrix:
        r = ranks(row)
        for i in range(n): R[i] += r[i]
        for t in Counter(row).values(): T += t ** 3 - t
    Rbar = m * (n + 1) / 2; S = sum((x - Rbar) ** 2 for x in R)
    den = m ** 2 * (n ** 3 - n) - m * T
    return 12 * S / den if den else float('nan')

def pct(xs, lo=0.025, hi=0.975):
    s = sorted(xs); return s[int(lo * (len(s) - 1))], s[int(math.ceil(hi * (len(s) - 1)))]

def sign_flip_p(d, rng, mc=20000):
    import itertools
    obs, n = abs(mean(d)), len(d)
    if n <= 16:
        sims = [abs(sum(x * s for x, s in zip(d, signs)) / n) for signs in itertools.product((1, -1), repeat=n)]
    else:
        sims = [abs(sum(x if rng.random() < .5 else -x for x in d) / n) for _ in range(mc)]
    return sum(s >= obs - 1e-12 for s in sims) / len(sims)

def holm(ps):
    order = sorted(range(len(ps)), key=lambda i: ps[i]); out, run = [0.0] * len(ps), 0.0
    for j, i in enumerate(order):
        run = max(run, min(1.0, (len(ps) - j) * ps[i])); out[i] = run
    return out

def f2(x, d=2):
    return '' if x is None or (isinstance(x, float) and math.isnan(x)) else (f'{x:.{d}f}' if isinstance(x, float) else str(x))

def md_table(head, rows):
    out = ['| ' + ' | '.join(head) + ' |', '|' + '---|' * len(head)]
    out += ['| ' + ' | '.join(str(c) for c in r) + ' |' for r in rows]
    return '\n'.join(out)