"""Score answers and summarise the benchmark. See PROTOCOL.md.

python3 analyze.py [results/answers.jsonl]
"""
import json
import random
import re
import sys
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).parent
CONDITIONS = ['A', 'B2', 'B4', 'C', 'D10', 'D25', 'Dall']
NAMES = {'A': 'Full text', 'B2': 'RAG, top 2 chunks', 'B4': 'RAG, top 4 chunks', 'C': 'Summary',
         'D10': 'Croutons, top 10', 'D25': 'Croutons, top 25', 'Dall': 'All croutons'}
# Published cached-input discounts (cached input price as a share of the normal input price).
CACHE_RATES = {'Anthropic (cache read)': 0.10, 'OpenAI (cached input)': 0.50, 'Google Gemini (implicit cache)': 0.25}
B = 10000
rng = random.Random(7)


def norm(s):
    s = s.lower()
    s = s.replace('–', '-').replace('—', '-').replace('’', "'").replace('‘', "'")
    return re.sub(r'\s+', ' ', s)


def correct(q, answer):
    a = norm(answer)
    abstained = 'not stated' in a
    if q.get('unanswerable'):
        return abstained
    if abstained:
        return False
    return any(all(norm(p.strip()) in a for p in g.split('&')) for g in q['gold'])


def mean(xs):
    return sum(xs) / len(xs) if xs else float('nan')


def boot(xs, f=mean):
    n = len(xs)
    vals = sorted(f([xs[rng.randrange(n)] for _ in range(n)]) for _ in range(B))
    return vals[int(0.025 * B)], vals[int(0.975 * B) - 1]


def boot_paired(pairs):
    """pairs: list of (x_cond, x_A); returns mean difference and 95% interval."""
    n = len(pairs)
    diffs = [a - b for a, b in pairs]
    vals = sorted(mean([diffs[rng.randrange(n)] for _ in range(n)]) for _ in range(B))
    return mean(diffs), vals[int(0.025 * B)], vals[int(0.975 * B) - 1]


def main(path):
    Q = {q['id']: q for q in json.load(open(ROOT / 'questions.json'))}
    rows = [json.loads(l) for l in open(path)]
    by = defaultdict(dict)
    for r in rows:
        r['correct'] = correct(Q[r['id']], r['answer'])
        by[r['cond']][r['id']] = r

    prep = [json.loads(l) for l in open(ROOT / 'prep' / 'prep_log.jsonl')] if (ROOT / 'prep' / 'prep_log.jsonl').exists() else []
    prep_tokens = defaultdict(lambda: defaultdict(int))
    for p in prep:
        prep_tokens[p['kind']][p['slug']] += p['prompt_tokens'] + p['output_tokens']

    base = by['A']
    out = {'conditions': {}, 'n_questions': len(base)}
    for c in CONDITIONS:
        rs = by.get(c)
        if not rs:
            continue
        ans = [r for r in rs.values() if not Q[r['id']].get('unanswerable')]
        una = [r for r in rs.values() if Q[r['id']].get('unanswerable')]
        acc = [1.0 if r['correct'] else 0.0 for r in ans]
        abst = [1.0 if r['correct'] else 0.0 for r in una]
        tok = [r['prompt_tokens'] for r in rs.values()]
        pre = [r['prefill_ms'] for r in rs.values()]
        d = {
            'name': NAMES[c],
            'n': len(rs),
            'input_tokens_mean': mean(tok), 'input_tokens_ci': boot(tok),
            'prefill_ms_mean': mean(pre),
            'accuracy': mean(acc), 'accuracy_ci': boot(acc) if acc else None,
            'abstention': mean(abst) if abst else None, 'abstention_ci': boot(abst) if abst else None,
            'made_up_on_unanswerable': (1 - mean(abst)) if abst else None,
        }
        if c != 'A' and base:
            pairs = [(1.0 if rs[i]['correct'] else 0.0, 1.0 if base[i]['correct'] else 0.0)
                     for i in rs if i in base and not Q[i].get('unanswerable')]
            if pairs:
                d['accuracy_diff_vs_A'] = boot_paired(pairs)
            d['token_reduction_vs_A'] = 1 - d['input_tokens_mean'] / mean([r['prompt_tokens'] for r in base.values()])
        kind = {'C': 'summary', 'D10': 'croutons', 'D25': 'croutons', 'Dall': 'croutons'}.get(c)
        if kind and prep_tokens[kind]:
            per_doc = mean(list(prep_tokens[kind].values()))
            d['prep_tokens_per_article'] = per_doc
            d['prep_tokens_per_question'] = per_doc / 10
            a_tok = mean([r['prompt_tokens'] for r in base.values()]) if base else None
            if a_tok and a_tok > d['input_tokens_mean']:
                d['break_even_questions'] = per_doc / (a_tok - d['input_tokens_mean'])
        out['conditions'][c] = d

    if base:
        a_tok = out['conditions']['A']['input_tokens_mean']
        out['caching_on_A'] = {k: {'effective_input_tokens_after_first_question': a_tok * v,
                                   'reduction': 1 - v} for k, v in CACHE_RATES.items()}

    h = {}
    if 'D25' in out['conditions'] and 'A' in out['conditions']:
        d25 = out['conditions']['D25']
        h['H1_token_reduction_ge_80pct'] = d25['token_reduction_vs_A'] >= 0.80
        h['H2_accuracy_within_5pts'] = d25.get('accuracy_diff_vs_A', (0, -1, 0))[1] > -0.05
    if ('D25' in out['conditions'] and 'B4' in out['conditions']
            and out['conditions']['D25']['made_up_on_unanswerable'] is not None
            and out['conditions']['B4']['made_up_on_unanswerable'] is not None):
        h['H3_made_up_le_B4'] = out['conditions']['D25']['made_up_on_unanswerable'] <= out['conditions']['B4']['made_up_on_unanswerable']
    out['hypotheses'] = h

    json.dump(out, open(ROOT / 'results' / ('summary_' + Path(path).stem + '.json'), 'w'), indent=1)
    print(f"{'Condition':22} {'tokens/q':>9} {'vs A':>7} {'accuracy':>16} {'diff vs A':>22} {'abstain':>8} {'prefill ms':>10}")
    for c, d in out['conditions'].items():
        ci = d['accuracy_ci']
        diff = d.get('accuracy_diff_vs_A')
        vs = '-' if c == 'A' else '%+.0f%%' % (-d['token_reduction_vs_A'] * 100)
        dtxt = '' if not diff else '%+5.1f [%+5.1f,%+5.1f]' % (diff[0] * 100, diff[1] * 100, diff[2] * 100)
        ab = '' if d['abstention'] is None else '%5.0f%%' % (d['abstention'] * 100)
        print('%-22s %9.0f %7s %5.1f [%4.1f,%5.1f] %22s %8s %10.0f' % (
            d['name'], d['input_tokens_mean'], vs, d['accuracy'] * 100, ci[0] * 100, ci[1] * 100, dtxt, ab, d['prefill_ms_mean']))
    for c in ('C', 'D25'):
        if c in out['conditions'] and 'prep_tokens_per_article' in out['conditions'][c]:
            d = out['conditions'][c]
            print(f"{c}: prep {d['prep_tokens_per_article']:.0f} tokens/article; break-even after {d.get('break_even_questions', float('nan')):.1f} questions")
    print('hypotheses:', h)


if __name__ == '__main__':
    main(sys.argv[1] if len(sys.argv) > 1 else str(ROOT / 'results' / 'answers.jsonl'))