"""
Can you trust an AI judge?  Step 3: the full data.
===================================================
Data: MT-Bench human judgments (lmsys/mt_bench_human_judgments, CC-BY-4.0),
from Zheng et al. 2023, "Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena".
3,355 votes by experts and paper authors on pairs of chatbot answers, and
GPT-4's verdict on the same pairs. GPT-4 judged every pair twice, with the
answers in both orders; when its two verdicts disagreed, the dataset records
"tie (inconsistent)".

Setup (once):  pip install pandas pyarrow
Run:           python judge_lesson.py            (full report)
               python judge_lesson.py read math  (read disagreements in one category)
The first run downloads the two data files (about 1.4 MB) into ./data.
"""
import os
import sys
import urllib.request
from math import sqrt
import pandas as pd

CATS = ['writing', 'roleplay', 'reasoning', 'math', 'coding', 'extraction', 'stem', 'humanities']
KEY = ['question_id', 'model_a', 'model_b', 'turn']


BASE = 'https://huggingface.co/datasets/lmsys/mt_bench_human_judgments/resolve/main/data/'
FILES = {'data/human.parquet': 'human-00000-of-00001-25f4910818759289.parquet',
         'data/gpt4_pair.parquet': 'gpt4_pair-00000-of-00001-c0b431264a82ddc0.parquet'}


def fetch():
    os.makedirs('data', exist_ok=True)
    for local, remote in FILES.items():
        if not os.path.exists(local):
            print('downloading', remote, '...')
            urllib.request.urlretrieve(BASE + remote, local)


def load():
    fetch()
    h = pd.read_parquet('data/human.parquet')
    g = pd.read_parquet('data/gpt4_pair.parquet')
    for d in (h, g):
        d['category'] = d.question_id.map(lambda q: CATS[(q - 81) // 10])
    return h, g


def wilson(k, n, z=1.96):
    """95% interval for a proportion k/n. Better than p +/- 1.96*sqrt(p(1-p)/n) near 0 or 1."""
    if n == 0:
        return (0.0, 1.0)
    p = k / n
    centre = (p + z * z / (2 * n)) / (1 + z * z / n)
    half = z * sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)
    return (centre - half, centre + half)


def judge_vs_human(h, g, cats=None):
    """Each decisive human vote compared with GPT-4's decisive verdict on the same pair."""
    m = h.merge(g[KEY + ['winner']].rename(columns={'winner': 'gpt4'}), on=KEY)
    if cats:
        m = m[m.category.isin(cats)]
    d = m[m.winner.isin(['model_a', 'model_b']) & m.gpt4.isin(['model_a', 'model_b'])]
    return int((d.winner == d.gpt4).sum()), len(d)


def human_vs_human(h, cats=None):
    """Every pair of decisive human votes on the same item."""
    hh = h if not cats else h[h.category.isin(cats)]
    k = n = 0
    for votes in hh[hh.winner != 'tie'].groupby(KEY).winner.agg(list):
        for i in range(len(votes)):
            for j in range(i + 1, len(votes)):
                n += 1
                k += votes[i] == votes[j]
    return k, n


def flip_rate(g, cats=None):
    gg = g if not cats else g[g.category.isin(cats)]
    return int((gg.winner == 'tie (inconsistent)').sum()), len(gg)


def fmt(k, n):
    lo, hi = wilson(k, n)
    return f"{k / n:5.0%}  ({lo:.0%} to {hi:.0%}, n={n})"


def report(h, g):
    print("Agreement, decisive votes only (ties left out)\n")
    print(f"{'category':<11} {'GPT-4 vs human':<27} {'human vs human':<27} {'GPT-4 flips':<12}")
    for c in CATS:
        k1, n1 = judge_vs_human(h, g, [c])
        k2, n2 = human_vs_human(h, [c])
        k3, n3 = flip_rate(g, [c])
        print(f"{c:<11} {fmt(k1, n1):<27} {fmt(k2, n2):<27} {k3 / n3:5.0%}")
    k1, n1 = judge_vs_human(h, g); k2, n2 = human_vs_human(h); k3, n3 = flip_rate(g)
    print(f"{'ALL':<11} {fmt(k1, n1):<27} {fmt(k2, n2):<27} {k3 / n3:5.0%}")

    # Same comparison, but ties count: three possible labels (A, B, tie) and
    # GPT-4's inconsistent verdicts are treated as ties.
    m = h.merge(g[KEY + ['winner']].rename(columns={'winner': 'gpt4'}), on=KEY)
    gl = m.gpt4.replace({'tie (inconsistent)': 'tie'})
    print(f"\nWith ties counted as a third label: GPT-4 vs human = {(m.winner == gl).mean():.0%} (n={len(m)})")


def read(h, g, cat, limit=4):
    """Show pairs where every human rater agreed and GPT-4 chose the other answer."""
    m = h[h.category == cat].groupby(KEY).agg(votes=('winner', list), conv_a=('conversation_a', 'first'),
                                              conv_b=('conversation_b', 'first')).reset_index()
    m = m.merge(g[KEY + ['winner']].rename(columns={'winner': 'gpt4'}), on=KEY)
    shown = 0
    for _, r in m.iterrows():
        v = set(r.votes)
        if len(v) == 1 and v <= {'model_a', 'model_b'} and r.gpt4 in ('model_a', 'model_b') and r.gpt4 not in v:
            t = r.turn - 1
            print("=" * 78)
            print(f"question {r.question_id}, turn {r.turn}: humans chose {r.votes[0][-1].upper()} "
                  f"({len(r.votes)} vote(s)), GPT-4 chose {r.gpt4[-1].upper()}")
            print("\nQ:", r.conv_a[2 * t]['content'][:700])
            print("\nA:", r.conv_a[2 * t + 1]['content'][:900])
            print("\nB:", r.conv_b[2 * t + 1]['content'][:900])
            shown += 1
            if shown == limit:
                break
    if not shown:
        print("No unanimous-human disagreements in", cat)


if __name__ == '__main__':
    h, g = load()
    if len(sys.argv) >= 3 and sys.argv[1] == 'read':
        read(h, g, sys.argv[2])
    else:
        report(h, g)