"""
Is the new model better, or is it noise?  Steps 4 to 6.
=======================================================
Data: MT-Bench single-answer grades. GPT-4 scored each model's answer to the
same 80 questions (two turns each, 160 items) on a 1 to 10 scale.
Source: Zheng et al. 2023, "Judging LLM-as-a-Judge with MT-Bench and Chatbot
Arena"; judgments published with LMSYS FastChat (huggingface.co/spaces/lmsys/mt-bench).

Setup (once):  pip install pandas numpy
Run:           python noise_lesson.py                              (the Llama-2 decision and the leaderboard)
               python noise_lesson.py compare gpt-3.5-turbo claude-v1   (any two models)
               python noise_lesson.py models                       (list the 34 models)
The first run downloads the grades (about 20 MB) into ./data.
"""
import os
import sys
import urllib.request

import numpy as np
import pandas as pd

BASE = 'https://huggingface.co/spaces/lmsys/mt-bench/resolve/main/data/mt_bench/'
FILES = {'data/gpt4_single.jsonl': 'model_judgment/gpt-4_single.jsonl'}
CATS = ['writing', 'roleplay', 'reasoning', 'math', 'coding', 'extraction', 'stem', 'humanities']
BOOT = 5000
rng = np.random.default_rng(0)


def load():
    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)
    d = pd.read_json('data/gpt4_single.jsonl', lines=True)
    d = d[d.score >= 0]
    d['category'] = d.question_id.map(lambda q: CATS[(q - 81) // 10])
    # one row per item (question, turn), one column per model
    return d.pivot_table(index=['question_id', 'turn', 'category'], columns='model',
                         values='score', aggfunc='first').reset_index()


def compare(w, a, b):
    x = w[['category', a, b]].dropna()
    diff = (x[a] - x[b]).to_numpy()
    n = len(diff)
    # Unpaired: treat the two score lists as unrelated samples.
    se_unpaired = np.sqrt(x[a].var(ddof=1) / n + x[b].var(ddof=1) / n)
    # Paired: both models answered the same questions, so look at the differences.
    se_paired = diff.std(ddof=1) / np.sqrt(n)
    boots = np.array([diff[rng.integers(0, n, n)].mean() for _ in range(BOOT)])
    print(f"\n{a}  vs  {b}   ({n} items)")
    print(f"  average score        {x[a].mean():.2f}  vs  {x[b].mean():.2f}")
    print(f"  difference           {diff.mean():+.2f}")
    print(f"  95% range, unpaired  +/- {1.96 * se_unpaired:.2f}")
    print(f"  95% range, paired    +/- {1.96 * se_paired:.2f}")
    print(f"  items: {a} higher {np.sum(diff > 0)}, lower {np.sum(diff < 0)}, same {np.sum(diff == 0)}")
    print(f"  resampled {BOOT} times: {a} came out ahead in {np.mean(boots > 0):.0%} of resamples")
    print("  by category (20 items each):")
    for c in CATS:
        dc = (x[x.category == c][a] - x[x.category == c][b]).to_numpy()
        print(f"    {c:<11} {dc.mean():+.2f}  +/- {1.96 * dc.std(ddof=1) / np.sqrt(len(dc)):.2f}")


def leaderboard(w, models):
    x = w[models].dropna()
    n = len(x)
    ranks = {m: [] for m in models}
    for _ in range(2000):
        r = x.iloc[rng.integers(0, n, n)].mean().rank(ascending=False, method='first')
        for m in models:
            ranks[m].append(r[m])
    print("\nThe leaderboard, and where each model lands when the questions are resampled")
    print(f"  {'model':<24} {'score':>5}   {'rank in 90% of resamples':<26} {'ranked 1st':>10}")
    for m in sorted(models, key=lambda m: -x[m].mean()):
        r = np.array(ranks[m])
        print(f"  {m:<24} {x[m].mean():5.2f}   {np.percentile(r, 5):.0f} to {np.percentile(r, 95):.0f}{'':<20} {np.mean(r == 1):>9.0%}")


if __name__ == '__main__':
    w = load()
    if len(sys.argv) == 4 and sys.argv[1] == 'compare':
        compare(w, sys.argv[2], sys.argv[3])
    elif len(sys.argv) == 2 and sys.argv[1] == 'models':
        print('\n'.join(sorted(c for c in w.columns if c not in ('question_id', 'turn', 'category'))))
    else:
        compare(w, 'Llama-2-70b-chat', 'Llama-2-13b-chat')
        leaderboard(w, ['tulu-30b', 'guanaco-65b', 'oasst-sft-7-llama-30b', 'palm-2-chat-bison-001',
                        'mpt-30b-chat', 'vicuna-13b-v1.3', 'wizardlm-13b'])