"""2.9 (English edition of zh 2.8_isolation): does processing many memos in one window make content bleed between them,
and does splitting them into one clean window each cost us cross-memo clues?

Batched: all 28 W+L transcripts in one prompt, cleaned up together (with names.md), output as JSON {id: text}.
Isolated: each memo cleaned up in its own call with the same instructions and the same names.md.
Name accuracy: the corpus scorer (corpus/score_names.py), 31 name mentions over the 28 W+L memos.
Crosstalk: a cast member's first name (whole word, case-insensitive) that appears in a cleaned memo whose script never
mentions that person. The Chinese check exempted the given name alone (思远 for 蔺思远); the English cast is referred to
by first name only, so the first name itself is the check. Pronoun references (W01 "she", W12 "her" = Saoirse) are in
memos that also name her, so they never trigger it.
temperature 0, 3 runs per condition, same design as the zh run.

Variants (all reported, see README):
  grok        main run: corpus/transcripts_grok.json, names.md with roles (the zh design)
  compat      same, on corpus/transcripts_compat.json
  norole      grok transcripts, names list without the role column
Conditions: batch, isolated (the zh design) and batch1 (English-only control: the batch prompt and JSON format with a
single memo per call; --conds batch1).

Usage: run.py [--variant grok] [--runs 3] [--models m1,m2] [--conds batch,isolated,batch1]   (existing out_*.json files are reused, not re-called)
       run.py --score            (re-score every out_*.json into results.json)
"""
import json, pathlib, re, sys
from concurrent.futures import ThreadPoolExecutor
HERE = pathlib.Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str(ROOT)); sys.path.insert(0, str(ROOT / 'corpus'))
from common import chat, dump, log_cost
from score_names import score, NAMES

d = json.load(open(ROOT / 'corpus/memos.json', encoding='utf-8'))
ids = [m['id'] for m in d['memos'] if m['set'] in 'WL']
ref = {m['id']: m['text'] for m in d['memos']}
PEOPLE = list(d['people'])
NAMES_MD = (ROOT / '2_3_names/names.md').read_text(encoding='utf-8')
NAMES_NOROLE = ('# People I mention often\n\n' + '\n'.join(f'- {p}' for p in PEOPLE) + '\n\nCommon terms: '
                + ', '.join(d['terms']) + '\n')
VARIANTS = {'grok': ('transcripts_grok.json', NAMES_MD),
            'compat': ('transcripts_compat.json', NAMES_MD),
            'norole': ('transcripts_grok.json', NAMES_NOROLE)}
ASK = 'Clean up these voice memo transcripts: fix recognition errors and punctuation, but don\'t rewrite the sentences.'
MODELS = ['deepseek-v4.1-flash', 'gemini-3.8-flash']
PARSE_RETRIES = []
word = lambda p: re.compile(rf'\b{re.escape(p)}\b', re.I)


def body(text):
    """The cleaned text only. A model sometimes appends a '---' note after the text (deepseek on W18: it leaves "Ify's"
    and then asks whether Aoife was meant, or says "Wait — ... the intern is Aoife" and repeats the text). Only the text
    above the note counts, because that is what would land in cleaned/; the note is recorded as a hedge."""
    return re.split(r'\n\s*(?:---|\*\*Note|Note:|Wait\b)', text)[0]


def crosstalk(outs):
    found = {}
    for mid in ids:
        extra = [p for p in PEOPLE if word(p).search(outs.get(mid, '')) and not word(p).search(ref[mid])]
        if extra:
            found[mid] = extra
    return found


def misses_by_memo(outs, src):
    """Which name mentions were missed, memo by memo, with the ASR form for context."""
    rows = []
    for mid in ids:
        for p in PEOPLE:
            n = len(word(p).findall(ref[mid]))
            if not n:
                continue
            k = min(n, len(word(p).findall(outs.get(mid, ''))))
            if k < n:
                rows.append({'memo': mid, 'name': p, 'missed': n - k, 'asr': src[mid], 'output': outs.get(mid, '')})
    return rows


def batch(model, src, names, only=None):
    sel = only or ids
    blob = '\n\n'.join(f'[{mid}] {src[mid]}' for mid in sel)
    n = f'{len(sel)} transcripts, each' if len(sel) > 1 else '1 transcript,'
    prompt = (f'{ASK}\n\n{names}\n\nBelow are {n} starting with [id]. Clean up each one and output '
              f'a single JSON object whose keys are the ids and whose values are the cleaned full texts. Output nothing else.\n\n{blob}')
    for attempt in range(3):  # an unparseable JSON reply is retried (recorded in parse_retries.json)
        out, u, _ = chat([{'role': 'user', 'content': prompt}], model, temperature=0, response_format={'type': 'json_object'})
        for cand in (out, (re.search(r'\{.*\}', out, re.S) or [''])[0]):
            try:
                return json.loads(cand), u
            except Exception:
                pass
        PARSE_RETRIES.append({'model': model, 'ids': sel, 'reply_head': out[:200]})
    return {}, u


def batch_of_one(model, src, names):
    """Control added in the English run (not in zh): the batched prompt and JSON output format, but one memo per call.
    If this scores like the batch, the batch gain comes from the prompt format, not from seeing the other memos."""
    def one(mid):
        o, _ = batch(model, src, names, only=[mid])
        return mid, (o.get(mid) or next(iter(o.values()), '')) if isinstance(o, dict) else ''
    with ThreadPoolExecutor(8) as ex:
        return dict(ex.map(one, ids)), None


def isolated(model, src, names):
    def one(mid):
        prompt = f'{ASK} Output only the cleaned full text.\n\n{names}\n\nTranscript:\n{src[mid]}'
        return mid, chat([{'role': 'user', 'content': prompt}], model, temperature=0)[0].strip()
    with ThreadPoolExecutor(8) as ex:
        return dict(ex.map(one, ids)), None


def score_all():
    res = {}
    for f in sorted(HERE.glob('out_*.json')):
        _, variant, model, cond, r = f.stem.split('_', 4)
        src = json.load(open(ROOT / 'corpus' / VARIANTS[variant][0], encoding='utf-8'))
        raw = json.load(open(f, encoding='utf-8'))
        outs = {k: body(v) for k, v in raw.items()}
        hedged = [k for k in raw if body(raw[k]) != raw[k]]
        tot, hit, miss, _ = score(outs, NAMES, ids=set(ids))
        ct = crosstalk(outs)
        mm = misses_by_memo(outs, src)
        key = f'{variant}:{model}:{cond}'
        res.setdefault(key, []).append({'run': int(r), 'names': f'{hit}/{tot}', 'names_correct': hit, 'crosstalk_memos': len(ct),
                                        'crosstalk': ct, 'hedged_with_note': hedged, 'missing_ids': [i for i in ids if i not in outs],
                                        'name_misses': miss, 'missed_mentions': [(x['memo'], x['name']) for x in mm]})
    for k, runs in res.items():
        runs.sort(key=lambda x: x['run'])
        print(f"{k:40s} names {[x['names'] for x in runs]}  crosstalk {[x['crosstalk_memos'] for x in runs]}  "
              f"missed {sorted(set(m for x in runs for m in map(tuple, x['missed_mentions'])))}")
    asr = {}
    for v in ('grok', 'compat'):
        src = json.load(open(ROOT / 'corpus' / VARIANTS[v][0], encoding='utf-8'))
        tot, hit, miss, _ = score(src, NAMES, ids=set(ids))
        asr[v] = {'names': f'{hit}/{tot}', 'crosstalk_memos_in_raw_asr': len(crosstalk(src))}
    summary = {k: {'names_per_run': [x['names'] for x in v], 'crosstalk_memos_per_run': [x['crosstalk_memos'] for x in v],
                   'names_from_another_memo': '0 in every run: each crosstalk hit (W20, and W14 in batch1) also occurs with one memo per call, so none came from another memo (see README)',
                   'hedged_per_run': [x['hedged_with_note'] for x in v]} for k, v in res.items()}
    dump({'summary': summary, 'asr_only': asr, 'batch_input_tokens': {'gemini_billed': 1959, 'cl100k_transcripts': 1682,
          'cl100k_names_md': 121}, 'runs': res}, HERE / 'results.json')


if __name__ == '__main__':
    args = sys.argv[1:]
    if '--score' in args:
        score_all(); sys.exit()
    variant = args[args.index('--variant') + 1] if '--variant' in args else 'grok'
    runs = int(args[args.index('--runs') + 1]) if '--runs' in args else 3
    models = args[args.index('--models') + 1].split(',') if '--models' in args else MODELS
    conds = args[args.index('--conds') + 1].split(',') if '--conds' in args else ['batch', 'isolated']
    fname, names = VARIANTS[variant]
    src = json.load(open(ROOT / 'corpus' / fname, encoding='utf-8'))
    usage = {}
    for model in models:
        for cond, fn in [('batch', batch), ('isolated', isolated), ('batch1', batch_of_one)]:
            if cond not in conds:
                continue
            for r in range(runs):
                f = HERE / f'out_{variant}_{model}_{cond}_{r}.json'
                if f.exists():
                    continue
                outs, u = fn(model, src, names)
                dump(outs, f)
                if u:
                    usage[f.stem] = u
                tot, hit, _, _ = score({k: body(v) for k, v in outs.items()}, NAMES, ids=set(ids))
                print(variant, model, cond, r, f'names {hit}/{tot}', 'crosstalk', crosstalk(outs), flush=True)
    if usage:
        uf = HERE / 'batch_usage.json'
        old = json.load(open(uf)) if uf.exists() else {}
        old.update(usage); dump(old, uf)
    if PARSE_RETRIES:
        pf = HERE / 'parse_retries.json'
        old = json.load(open(pf)) if pf.exists() else []
        dump(old + PARSE_RETRIES, pf)
    log_cost(f'2_9_isolation/run.py --variant {variant} --runs {runs}')
    score_all()