"""
Do your human raters agree?  Steps 3 to 6.
==========================================
Data: HelpSteer2 individual annotations (nvidia/HelpSteer2, "disagreements",
CC-BY-4.0). 23,652 chatbot responses. Trained annotators rated each response
from 0 to 4 on helpfulness, correctness, coherence, complexity and verbosity.
Every response got at least 3 annotators; if the first ones disagreed by more
than 2 points on helpfulness, up to 2 more were added.
Wang et al., 2024, "HelpSteer2: Open-source dataset for training top-performing reward models."

Setup (once):  pip install pandas numpy
Run:           python raters_lesson.py              (agreement for each attribute)
               python raters_lesson.py pooling      (a common way to get these numbers wrong)
               python raters_lesson.py raters       (how many raters per item?)
               python raters_lesson.py read         (read responses where raters split 0 to 4)
The first run downloads about 13 MB into ./data.
"""
import os
import sys
import urllib.request

import numpy as np
import pandas as pd

URL = 'https://huggingface.co/datasets/nvidia/HelpSteer2/resolve/main/disagreements/disagreements.jsonl.gz'
ATTRS = ['helpfulness', 'correctness', 'coherence', 'complexity', 'verbosity']
SKIP = ('weapon', 'bomb', 'explosive', 'nuclear', 'kill', 'sex', 'porn', 'drug', 'suicide', 'nigg', 'rape', 'eradicat', 'genocid', 'ethnic', 'amalek', 'terror', 'nazi')


def load():
    os.makedirs('data', exist_ok=True)
    if not os.path.exists('data/hs2_disagreements.jsonl.gz'):
        print('downloading HelpSteer2 annotations ...')
        urllib.request.urlretrieve(URL, 'data/hs2_disagreements.jsonl.gz')
    return pd.read_json('data/hs2_disagreements.jsonl.gz', lines=True)


def one_pair(ratings, seed=0):
    """Pick two raters at random for every response, so each response counts once."""
    rng = np.random.default_rng(seed)
    return np.array([rng.choice(r, 2, replace=False) for r in ratings])


def all_pairs(ratings):
    """Every pair of raters on every response. A response with 5 raters gives 10 pairs, one with 3 gives 3."""
    return np.array([(r[i], r[j]) for r in ratings for i in range(len(r)) for j in range(i + 1, len(r))])


def agreement(pairs):
    a, b = pairs[:, 0], pairs[:, 1]
    exact = np.mean(a == b)
    within1 = np.mean(np.abs(a - b) <= 1)
    share = np.bincount(pairs.ravel(), minlength=5) / pairs.size
    chance = np.sum(share ** 2)                     # two raters picking the same score by luck
    kappa = (exact - chance) / (1 - chance)         # agreement beyond chance, 0 = luck, 1 = perfect
    near = 1 - (np.arange(5)[:, None] - np.arange(5)[None, :]) ** 2 / 16   # partial credit for near misses
    weighted = (np.mean(near[a, b]) - np.sum(np.outer(share, share) * near)) / (1 - np.sum(np.outer(share, share) * near))
    return exact, within1, chance, kappa, weighted, share


def report(d):
    print("Two randomly chosen raters per response, all 23,652 responses\n")
    print(f"  {'attribute':<12} {'exact':>6} {'within 1':>9} {'by chance':>10} {'kappa':>6} {'weighted':>9}   share of each score 0..4")
    for attr in ATTRS:
        ex, w1, ch, k, wk, share = agreement(one_pair(d[attr]))
        print(f"  {attr:<12} {ex:>6.0%} {w1:>9.0%} {ch:>10.0%} {k:>6.2f} {wk:>9.2f}   {' '.join(f'{s:.2f}' for s in share)}")
    p = one_pair(d.helpfulness) >= 3
    ex = np.mean(p[:, 0] == p[:, 1]); pr = p.mean(); ch = pr ** 2 + (1 - pr) ** 2
    print(f"\n  helpfulness as yes/no (3 or 4 = yes): exact {ex:.0%}, by chance {ch:.0%}, kappa {(ex - ch) / (1 - ch):.2f}")


def pooling(d):
    n = d.helpfulness.map(len)
    print("Helpfulness agreement computed three ways\n")
    for label, pairs in [('one pair per response', one_pair(d.helpfulness)),
                         ('every pair pooled', all_pairs(d.helpfulness))]:
        ex, w1, ch, k, wk, _ = agreement(pairs)
        print(f"  {label:<24} pairs {len(pairs):>7,}   exact {ex:.0%}   kappa {k:.2f}   weighted {wk:.2f}")
    print("\nBy how many raters a response got")
    for label, sub in [('3 raters', d[n == 3]), ('5 or more raters', d[n >= 5])]:
        ex, w1, ch, k, wk, _ = agreement(one_pair(sub.helpfulness))
        r = np.corrcoef(one_pair(sub.helpfulness).T)[0, 1]
        print(f"  {label:<18} {len(sub):>7,} responses   exact {ex:.0%}   kappa {k:.2f}   correlation {r:.2f}")


def raters(d):
    p = one_pair(d.helpfulness)
    r1 = np.corrcoef(p.T)[0, 1]
    print(f"Correlation between two single raters on helpfulness: {r1:.2f}")
    print("Reliability of the AVERAGE of k raters (Spearman-Brown: k*r / (1 + (k-1)*r)):\n")
    for k in (1, 2, 3, 5, 10):
        print(f"  {k:>2} raters   {k * r1 / (1 + (k - 1) * r1):.2f}")


def read(d, limit=4):
    d = d.copy()
    d['spread'] = d.helpfulness.map(lambda x: max(x) - min(x))
    short = d[(d.prompt.str.split().str.len() <= 60) & (d.response.str.split().str.len().between(15, 150))
              & ~d.prompt.str.contains('<extra_id_1>')
              & ~d.prompt.str.lower().str.contains('|'.join(SKIP)) & ~d.response.str.lower().str.contains('|'.join(SKIP))]
    split = short[short.spread == 4].sample(limit, random_state=11)
    for _, r in split.iterrows():
        print("=" * 78)
        print(f"helpfulness ratings {r.helpfulness}   correctness {r.correctness}")
        print("\nPROMPT:", r.prompt)
        print("\nRESPONSE:", r.response)


if __name__ == '__main__':
    d = load()
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    {'pooling': pooling, 'raters': raters, 'read': read}.get(cmd, report)(d)