"""2.4 (English): the same cleanup request in a fresh session, with and without the cleanup skill.
Twin of experiments/2_4_skill/run.py. Both conditions get the names list (names.md), so the only difference is the skill.
Input: the raw Grok transcripts (corpus/transcripts_grok.json) of the 28 W+L memos, one fresh call per memo, temperature 0.

Checks per memo (all three must hold for a first-try pass; the Chinese run had four, but the English ASR already strips
um/uh, so "fillers removed" has nothing to measure):
  keep_tic   : "I mean" kept as many times as it appears in the transcript the model was given
               (zh counted against the script, which produced 3 artefactual failures where ASR had dropped a word)
  numbers    : house style for numbers. In the 5 memos where the transcript spells a number out (W02 "ten in the morning",
               W08 "eight at night", W12 "ten minute demo", L02 "around nine", L05 "forty cents") the digit form must be
               there and the word gone, with times written "10 a.m." / "8 p.m."; in every memo, no time in another
               format ("10 AM", "10am", "10 a.m" ...).
  no_rewrite : output word count within 85%-115% of the transcript's (nothing summarised or merged away)
Diagnostics (not part of the pass): names right per names.md, preamble/markdown, contractions expanded, casual markers
dropped, word-level similarity to the transcript.
Usage: run.py [model ...]   (default: deepseek-v4.1-flash gemini-3.8-flash)"""
import difflib
import json
import pathlib
import re
import sys
from concurrent.futures import ThreadPoolExecutor

ROOT = pathlib.Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT))
from common import chat, dump, log_cost

HERE = pathlib.Path(__file__).resolve().parent
d = json.load(open(ROOT / 'corpus/memos.json', encoding='utf-8'))
src = json.load(open(ROOT / 'corpus/transcripts_grok.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']}
names = (HERE / 'names.md').read_text(encoding='utf-8')
skill = (HERE / 'cleanup_skill.md').read_text(encoding='utf-8')
PEOPLE = list(d['people'])

TIC = re.compile(r'\bI mean\b', re.I)
# memo -> (pattern that must appear, spelled word that must be gone)
NUM_TARGETS = {
    'W02': (r'\b10(:00)? a\.m\.', r'\bten\b'),
    'W08': (r'\b8(:00)? p\.m\.', r'\beight\b'),
    'W12': (r'\b10[- ]minute', r'\bten\b'),
    'L02': (r'\b9(:00)?\b', r'\bnine\b'),
    'L05': (r'\b40 ?(cents|¢)|\$0?\.40\b', r'\bforty\b'),
}
# memos where at least one rule has something to check: the 11 with "I mean" plus the 5 number targets (L05 has both)
APPLICABLE = sorted({m for m in ids if TIC.search(src[m])} | set(NUM_TARGETS))
BAD_TIME = re.compile(r'\b\d{1,2}(:\d\d)?\s?(AM|PM|am|pm|A\.M\.|P\.M\.)(?![\w.])|\b\d{1,2}(:\d\d)?\s?[ap]\.m(?!\.)')
CASUAL = ['so', 'honestly', 'just', 'really', 'way too', 'pretty', 'kind of', 'basically', 'like', 'yeah', 'anyway',
          'note to self', 'got', 'gonna', 'stuff', 'thing']
CONTRACTION = re.compile(r"\b\w+(?:'|’)(?:s|re|ll|ve|d|t|m)\b", re.I)


def words(s):
    return re.findall(r"[A-Za-z0-9]+(?:['’.-][A-Za-z0-9]+)*", s)


def count_phrase(p, s):
    return len(re.findall(rf'\b{re.escape(p)}\b', s, re.I))


def check(mid, out):
    t = src[mid]
    num_ok = not BAD_TIME.search(out)
    if mid in NUM_TARGETS:
        need, gone = NUM_TARGETS[mid]
        num_ok = num_ok and bool(re.search(need, out)) and not re.search(gone, out, re.I)
    return {'keep_tic': len(TIC.findall(out)) >= len(TIC.findall(t)),
            'numbers': num_ok,
            'no_rewrite': 0.85 <= len(words(out)) / max(1, len(words(t))) <= 1.15}


def diagnose(mid, out):
    t = src[mid]
    name_tot = name_hit = 0
    for n in PEOPLE:
        k = len(re.findall(rf'\b{n}\b', ref[mid], re.I))
        if k:
            name_tot += k
            name_hit += min(k, len(re.findall(rf'\b{n}\b', out, re.I)))
    casual_t = sum(count_phrase(c, t) for c in CASUAL)
    casual_o = sum(min(count_phrase(c, t), count_phrase(c, out)) for c in CASUAL)
    return {'names': [name_hit, name_tot],
            'preamble_or_markdown': bool(re.match(r'\s*(here|sure|cleaned|okay|ok\b)', out, re.I)) or bool(re.search(r'\*\*|^#|^\s*[-*] ', out, re.M)),
            'contractions': [len(CONTRACTION.findall(t)), len(CONTRACTION.findall(out))],
            'casual_markers': [casual_t, casual_o],
            'tic': [len(TIC.findall(t)), len(TIC.findall(out))],
            'word_ratio': round(len(words(out)) / max(1, len(words(t))), 3),
            'similarity': round(difflib.SequenceMatcher(None, [w.lower() for w in words(t)], [w.lower() for w in words(out)]).ratio(), 3)}


ASK = 'Clean up this voice memo transcript. Output only the cleaned-up full text.'
conds = {'no_skill': lambda t: f'{ASK}\n\n{names}\n\nTranscript:\n{t}',
         'with_skill': lambda t: f'{ASK}\n\n{skill}\n\n{names}\n\nTranscript:\n{t}'}


def summarize(model, c, outs):
    checks = {mid: check(mid, outs[mid]) for mid in ids}
    diags = {mid: diagnose(mid, outs[mid]) for mid in ids}
    per_rule = {k: sum(v[k] for v in checks.values()) for k in next(iter(checks.values()))}
    tic_ids = [m for m in ids if TIC.search(src[m])]
    failed = {k: [m for m in ids if not checks[m][k]] for k in per_rule}
    return {
        'first_try_pass': f'{sum(all(v.values()) for v in checks.values())}/{len(ids)}',
        'per_rule': per_rule,
        'per_rule_applicable': {'keep_tic': f'{sum(checks[m]["keep_tic"] for m in tic_ids)}/{len(tic_ids)}',
                                'numbers': f'{sum(checks[m]["numbers"] for m in NUM_TARGETS)}/{len(NUM_TARGETS)} target memos',
                                'no_rewrite': f'{per_rule["no_rewrite"]}/{len(ids)}'},
        'pass_on_memos_with_something_to_check': f"{sum(all(checks[m].values()) for m in APPLICABLE)}/{len(APPLICABLE)}",
        'failed_memos': failed,
        'violations': sum(len(ids) - x for x in per_rule.values()),
        'diagnostics': {
            'names_right': f'{sum(v["names"][0] for v in diags.values())}/{sum(v["names"][1] for v in diags.values())}',
            'i_mean_kept': f'{sum(min(*v["tic"]) for v in diags.values())}/{sum(v["tic"][0] for v in diags.values())}',
            'memos_with_preamble_or_markdown': sum(v['preamble_or_markdown'] for v in diags.values()),
            'contractions_in_out': f'{sum(v["contractions"][1] for v in diags.values())} (transcripts: {sum(v["contractions"][0] for v in diags.values())})',
            'casual_markers_kept': f'{sum(v["casual_markers"][1] for v in diags.values())}/{sum(v["casual_markers"][0] for v in diags.values())}',
            'mean_word_ratio': round(sum(v['word_ratio'] for v in diags.values()) / len(ids), 3),
            'mean_similarity': round(sum(v['similarity'] for v in diags.values()) / len(ids), 3),
        },
        'per_memo': {m: {**checks[m], **diags[m]} for m in ids},
    }


if __name__ == '__main__':
    rescore = '--rescore' in sys.argv
    res = json.load(open(HERE / 'results.json', encoding='utf-8')) if (HERE / 'results.json').exists() else {}
    for model in [a for a in sys.argv[1:] if not a.startswith('--')] or ['deepseek-v4.1-flash', 'gemini-3.8-flash']:
        for c, mk in conds.items():
            f = HERE / f'out_{model}_{c}.json'
            if rescore:
                outs = json.load(open(f, encoding='utf-8'))
            else:
                def one(mid):
                    out, _, _ = chat([{'role': 'user', 'content': mk(src[mid])}], model, temperature=0)
                    return mid, out.strip()
                with ThreadPoolExecutor(8) as ex:
                    outs = dict(ex.map(one, ids))
                dump(outs, f)
            res[f'{model}:{c}'] = summarize(model, c, outs)
            r = res[f'{model}:{c}']
            print(model, c, r['first_try_pass'], r['per_rule'], r['per_rule_applicable'], r['diagnostics'])
            print('   failed:', r['failed_memos'])
    res['_applicable_memos'] = APPLICABLE
    dump(res, HERE / 'results.json')
    if not rescore:
        log_cost('2_4_skill/run.py', note='28 W+L memos x 2 conditions x 2 models')