"""
Did the model make that up?  Steps 3 to 6.
==========================================
Data: RAGTruth (Niu et al., 2024; github.com/ParticleMedia/RAGTruth, MIT license).
2,965 tasks, each answered by six models (17,790 answers). Every model was given
source documents and asked to use them: answer a question from web passages (QA),
summarize a news article (Summary), or describe a business from its listing data
(Data2txt). People then marked, word by word, anything in each answer that the
sources don't support or that contradicts them.

Setup (once):  pip install pandas
Run:           python madeup_lesson.py              (how often each model made things up)
               python madeup_lesson.py detect       (test two cheap automatic checks)
               python madeup_lesson.py nulls        (a common mistake with business data)
The first run downloads about 37 MB into ./data.
"""
import json
import os
import re
import sys
import urllib.request
from math import sqrt

import pandas as pd

BASE = 'https://raw.githubusercontent.com/ParticleMedia/RAGTruth/main/dataset/'
MODELS = ['gpt-3.5-turbo-0613', 'gpt-4-0613', 'llama-2-70b-chat', 'llama-2-13b-chat', 'llama-2-7b-chat', 'mistral-7B-instruct']
# Capitalized words that start sentences or are too common to count as names
COMMON = set("The A An In On At It Its This That These Those He She They We You I Here There However Yes No Overall "
             "Based According Some Many Also And But Or If When While For With As By To Of From Is Are Was Were Be "
             "Additionally Furthermore Finally First Second Third Unfortunately Therefore Passage Passages".split())


def load():
    os.makedirs('data', exist_ok=True)
    for name in ('response', 'source_info'):
        path = f'data/ragtruth_{name}.jsonl'
        if not os.path.exists(path):
            print('downloading', name, '...')
            urllib.request.urlretrieve(BASE + name + '.jsonl', path)
    r = pd.read_json('data/ragtruth_response.jsonl', lines=True)
    s = pd.read_json('data/ragtruth_source_info.jsonl', lines=True)
    m = r.merge(s[['source_id', 'task_type', 'source_info']], on='source_id')
    # An answer is "made up" if people marked at least one span the sources don't support.
    # (Annotators also marked some spans as true-but-unsupported with implicit_true; those don't count.)
    m['madeup'] = m.labels.map(lambda spans: any(not x.get('implicit_true', False) for x in spans))
    m['source'] = m.source_info.map(lambda x: x if isinstance(x, str) else json.dumps(x))
    return m


def wilson(k, n, z=1.96):
    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 rates(m):
    print("Share of answers with at least one made-up or contradicting passage (95% range)\n")
    print(f"  {'model':<22} {'all tasks':>18}   {'QA':>6} {'Summary':>8} {'Data2txt':>9}")
    for model in MODELS:
        x = m[m.model == model]
        lo, hi = wilson(x.madeup.sum(), len(x))
        by = x.groupby('task_type').madeup.mean()
        print(f"  {model:<22} {x.madeup.mean():>5.0%} ({lo:.0%} to {hi:.0%})   {by['QA']:>6.0%} {by['Summary']:>8.0%} {by['Data2txt']:>9.0%}")
    print(f"\n  {'all models':<22} {m.madeup.mean():>5.0%}")


def unsupported_numbers(answer, source):
    src = source.lower().replace(',', '')
    for n in re.findall(r"\b\d[\d,.]*\b", answer):
        n = n.strip('.,').replace(',', '')
        if n and n not in src:
            return True
    return False


def unsupported_names(answer, source):
    src = source.lower()
    for name in re.findall(r"\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b", answer):
        if name not in COMMON and name.lower() not in src:
            return True
    return False


def detect(m):
    checks = {
        'numbers only': lambda row: unsupported_numbers(row.response, row.source),
        'numbers or names': lambda row: unsupported_numbers(row.response, row.source) or unsupported_names(row.response, row.source),
    }
    print("Flag an answer if it contains a number (or a capitalized name) that doesn't appear in its sources.\n")
    print(f"  {'check':<18} {'answers':<22} {'flagged':>8} {'caught (recall)':>16} {'flags that were real (precision)':>33}")
    for label, check in checks.items():
        flag = m.apply(check, axis=1)
        for group, mask in [('all 17,790', m.model.notna()), ('GPT-4 only', m.model == 'gpt-4-0613'),
                            ('Mistral only', m.model == 'mistral-7B-instruct')]:
            f, y = flag[mask], m.madeup[mask]
            tp = (f & y).sum()
            print(f"  {label:<18} {group:<22} {f.mean():>8.0%} {tp / y.sum():>16.0%} {tp / f.sum():>33.0%}")
        print()
    spans = [x for spans in m.labels for x in spans if not x.get('implicit_true', False)]
    share = sum(bool(re.search(r"\d", x['text'])) for x in spans) / len(spans)
    print(f"Of {len(spans):,} marked passages, {share:.0%} contain a digit.")


def nulls(m):
    x = m[m.task_type == 'Data2txt']
    spans = [s for spans in x.labels for s in spans if not s.get('implicit_true', False)]
    null = [s for s in spans if 'null' in (s.get('meta') or '').lower()]
    print(f"Business-description answers: {len(spans):,} marked passages; {len(null):,} ({len(null) / len(spans):.0%}) have an annotator note mentioning a null (empty) field.\n")
    print("A few of them, with the annotator's note:")
    for s in null[:6]:
        print(f"  - \"{s['text'].strip()}\"   note: {(s.get('meta') or '').replace(chr(10), ' ')[:110]}")


if __name__ == '__main__':
    m = load()
    cmd = sys.argv[1] if len(sys.argv) > 1 else ''
    {'detect': detect, 'nulls': nulls}.get(cmd, rates)(m)