"""Recompute the frozen improvement trial locally, without model/API calls."""
from __future__ import annotations

import argparse
from collections import Counter
import hashlib
import json
from pathlib import Path
import re
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from app.engine.summary_consensus import select_consensus


def load_rows(path):
    rows = [json.loads(line) for line in path.read_text().splitlines() if line.strip()]
    assert len(rows) == 1000 and {r['id'] for r in rows} == set(range(1, 1001)), path
    return {r['id']: r for r in rows}


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('directory', type=Path)
    args = ap.parse_args()
    p = args.directory
    manifest = json.loads((p/'dataset_manifest.json').read_text())
    assert hashlib.sha256((p/'sources.jsonl').read_bytes()).hexdigest() == manifest['sha256']
    sources = load_rows(p/'sources.jsonl')
    protocol = json.loads((p/'protocol.json').read_text())
    assert not ({r['source_file'] for r in sources.values()} & set(protocol['development_books']))
    refs = load_rows(p/'reference.jsonl')
    systems = load_rows(p/'system.jsonl')
    baseline = load_rows(p/'baseline.jsonl')
    result = json.loads((p/'result.json').read_text())
    for label, outputs in [('baseline', baseline), ('improved', systems)]:
        for mode in ('character', 'word'):
            totals = {'precision': 0.0, 'recall': 0.0, 'f1': 0.0}
            for rid in range(1, 1001):
                candidates = [c['summary'] for c in outputs[rid]['candidates']]
                chosen = select_consensus(candidates) if label == 'improved' and result['selected'] == 'consensus5' else 0
                hyp = candidates[chosen] if chosen >= 0 else ''
                ref = refs[rid]['candidates'][0]['summary']
                def grams(text):
                    tokens = list(re.sub(r'\s+', '', text)) if mode == 'character' else text.split()
                    return Counter(zip(tokens, tokens[1:]))
                a, b = grams(ref), grams(hyp)
                overlap = sum(min(count, b[g]) for g, count in a.items())
                a_total, b_total = sum(a.values()), sum(b.values())
                totals['precision'] += overlap / b_total if b_total else 0.0
                totals['recall'] += overlap / a_total if a_total else 0.0
                totals['f1'] += 2 * overlap / (a_total + b_total) if a_total + b_total else 0.0
            actual = {key: value/1000 for key, value in totals.items()}
            assert all(abs(actual[key]-result[label][mode][key]) < 1e-12 for key in actual)
            print(label, mode, 'F1=%.6f%%' % (actual['f1']*100))
    print('PASS: 1000 IDs, source SHA-256, development-book exclusion, all P/R/F1 values')


if __name__ == '__main__':
    main()