"""2.6 (English): one correct edit becomes a wrong global rule, and how batch-history distillation fixes it.
Twin of experiments/2_5_overgeneralize/run.py, with the homonym "shell" (see PILOT.md) in place of 苹果.

The author writes the command-line sense of "shell" as "terminal" (the word on his screen). In W06 he makes one edit:
"run intake from the shell" -> "run intake from the terminal".
Step 1 (single-edit distillation): show only W06 before/after and ask for one rule for the cleanup rule file.
        5 runs per model, temperature 1. Scoped vs global is judged by hand (see README); a keyword flag is stored too.
Step 2 (downstream): clean up the other memos with (a) no rule, (b) the global rule, (c) a scoped rule, temperature 0.
        Work positives W03/W13/W20 (3 "shell" = command line; wanted: -> "terminal").
        Life negatives L01-L08 (9 "shell": gas station, beach, eggshell, taco x2, gas station, pasta, pistachio, hermit
        crab; unwanted: -> "terminal"). Undo = condition (a). A life memo where the global rule was not applied is
        "resisted"; outputs that comment on the rule are flagged separately.
Step 3 (batch-history distillation): retrieve all 12 memos in the 38-memo history that contain "shell": 4 work memos
        (W03/W06/W13/W20) where it was edited to "terminal", 8 life memos (L01-L08) left unchanged. Ask for one rule from
        the batch, 5 runs per model, temperature 1.
Input: corpus/transcripts_grok.json. Usage: run.py [--batch-only] [model ...];  run.py --rescore (no API calls)"""
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
src = json.load(open(ROOT / 'corpus/transcripts_grok.json', encoding='utf-8'))
names = (ROOT / '2_4_skill/names.md').read_text(encoding='utf-8')
SHELL = re.compile(r'\bshell\b', re.I)
TERM = re.compile(r'\bterminals?\b', re.I)


def edit(t):
    return SHELL.sub(lambda m: 'Terminal' if m.group(0)[0].isupper() else 'terminal', t)


before = src['W06']; after = edit(before)
PROPOSE = ('I made one manual edit while cleaning up a voice memo transcript. Before:\n' + before + '\n\nAfter:\n' + after +
           '\n\nTurn this edit into one rule for my cleanup rule file, so it is applied automatically whenever I clean up '
           'a transcript. Output only the rule, in one sentence.')
TECH, LIFE = ['W03', 'W13', 'W20'], ['L01', 'L02', 'L03', 'L04', 'L05', 'L06', 'L07', 'L08']
BATCH_IDS = ['W03', 'W06', 'W13', 'W20'] + LIFE
batch_records = [f"[{mid}]\nBefore: {src[mid]}\nAfter: {edit(src[mid]) if mid.startswith('W') else src[mid]}"
                 for mid in BATCH_IDS]
PROPOSE_BATCH = ('I searched my history for "shell" and found all 12 cleanup records that contain it (some I edited, '
                 'some I left alone):\n\n' + '\n\n'.join(batch_records) +
                 '\n\nBased on this batch of records, write one cleanup rule about "shell" for my cleanup rule file, so it '
                 'is applied automatically whenever I clean up a transcript. Output only the rule, in one sentence.')
RULES = {'none': '',
         'global': 'Rule: in transcripts, always write "shell" as "terminal".',
         'scoped': 'Rule: write "shell" as "terminal" only when it means the command line on a computer; '
                   'in any everyday sense, keep "shell".'}
ASK = 'Clean up this voice memo transcript: fix recognition errors and output only the cleaned-up full text.'
SCOPE_WORDS = re.compile(r'\b(only|when|whenever|unless|except|context|refer\w*|mean\w*|command|computing|technical|'
                         r'software|scripts?|code|cli|programming|keep|leave|everyday|unless)\b', re.I)


# Each "shell" mention is located by the words around it, so a model's commentary ("per your rule, shell -> terminal")
# is never counted as a conversion. The captured word is classified as shell / terminal / other / missing.
ANCHORS = {
    'W03': [r'both from (?:the|a) (\w+)'], 'W13': [r'working in (?:the|a) (\w+)'], 'W20': [r'into (?:the|a) (\w+) and rerun'],
    'L01': [r'stop at (?:the|a) (\w+)'], 'L02': [r'find (?:a|one) (\w+) to take'], 'L03': [r'no (\w+) left'],
    'L04': [r'warm each (\w+)', r'fill (?:the|each) (\w+)'], 'L05': [r'at (?:the|a) (\w+) by the school'],
    'L06': [r'that (\w+) pasta'], 'L07': [r'still in (?:the|their) (\w+)'], 'L08': [r'into a new (\w+)'],
}


def classify(out, mid):
    res = []
    for a in ANCHORS[mid]:
        m = re.search(a, out, re.I)
        w = m.group(1).lower() if m else None
        res.append('missing' if w is None else 'shell' if w.startswith('shell') else 'terminal' if w.startswith('terminal') else w)
    return res


def score(outs, src_words):
    cls = {mid: classify(outs[mid], mid) for mid in TECH + LIFE}
    flat = lambda ids, v: sum(c == v for i in ids for c in cls[i])
    commentary = [mid for mid in TECH + LIFE
                  if len(outs[mid].split()) > 1.25 * src_words[mid] or re.search(r'\*\*|cleaned-up text|per your rule', outs[mid], re.I)]
    return {'work_to_terminal': f"{flat(TECH, 'terminal')}/3", 'work_kept_shell': f"{flat(TECH, 'shell')}/3",
            'life_wrongly_terminal': f"{flat(LIFE, 'terminal')}/9", 'life_kept_shell': f"{flat(LIFE, 'shell')}/9",
            'life_memos_not_converted': [i for i in LIFE if 'terminal' not in cls[i]],
            'memos_with_commentary': commentary, 'per_mention': cls}


def downstream(model, rule):
    def one(mid):
        out, _, _ = chat([{'role': 'user', 'content': f'{ASK}\n\n{names}\n{rule}\n\nTranscript:\n{src[mid]}'}], model,
                         temperature=0)
        return mid, out.strip()
    with ThreadPoolExecutor(8) as ex:
        outs = dict(ex.map(one, TECH + LIFE))
    return {**score(outs, SRC_WORDS), 'outputs': outs}


SRC_WORDS = {k: len(v.split()) for k, v in src.items()}


if __name__ == '__main__':
    batch_only = '--batch-only' in sys.argv
    if '--rescore' in sys.argv:
        res = json.load(open(HERE / 'results.json', encoding='utf-8'))
        for k, v in res['downstream'].items():
            res['downstream'][k] = {**score(v['outputs'], SRC_WORDS), 'outputs': v['outputs']}
            print(k, {kk: vv for kk, vv in res['downstream'][k].items() if kk not in ('outputs', 'per_mention')})
        dump(res, HERE / 'results.json')
        sys.exit()
    models = [a for a in sys.argv[1:] if not a.startswith('--')] or ['deepseek-v4.1-flash', 'gemini-3.8-flash']
    res = json.load(open(HERE / 'results.json', encoding='utf-8')) if (HERE / 'results.json').exists() else {}
    res['edit'] = {'memo': 'W06', 'before': before, 'after': after}
    for k in ('proposals', 'batch_proposals', 'downstream'):
        res.setdefault(k, {})
    for model in models:
        if not batch_only:
            with ThreadPoolExecutor(5) as ex:
                props = list(ex.map(lambda _: chat([{'role': 'user', 'content': PROPOSE}], model, temperature=1)[0].strip(), range(5)))
            res['proposals'][model] = [{'rule': p, 'scope_keyword_flag': bool(SCOPE_WORDS.search(p))} for p in props]
            print(model, 'single:'); [print('   ', p) for p in props]
        with ThreadPoolExecutor(5) as ex:
            bprops = list(ex.map(lambda _: chat([{'role': 'user', 'content': PROPOSE_BATCH}], model, temperature=1)[0].strip(), range(5)))
        res['batch_proposals'][model] = [{'rule': p, 'scope_keyword_flag': bool(SCOPE_WORDS.search(p))} for p in bprops]
        print(model, 'batch:'); [print('   ', p) for p in bprops]
        if not batch_only:
            for rname, rule in RULES.items():
                r = downstream(model, rule)
                res['downstream'][f'{model}:{rname}'] = r
                print(model, rname, {k: v for k, v in r.items() if k != 'outputs'})
    dump(res, HERE / 'results.json')
    log_cost('2_6_overgeneralize/run.py' + (' --batch-only' if batch_only else ''),
             note='steps 1-3, 5 runs per distillation, 11 memos x 3 rules downstream')