"""2.5: one correct edit becomes a wrong global rule — and how batch history distillation fixes it.

Step 1a (single-edit rule extraction): the author edits one transcript (W06: 苹果的语音备忘录 -> Apple 的语音备忘录, he writes the
company name in English). We ask a model to turn that single edit into a rule for the memory file, 5 times per model,
and record whether the proposed rule is global or scoped.
Step 1b (batch-history rule extraction): instead of summarizing from a single edit on the spot, we search the historical
archive for all 12 memos containing 苹果 (4 work memos W03/W06/W13/W20 where 苹果 was edited to Apple, and 8 life memos
L01-L08 where 苹果 meant fruit and was left unchanged) and ask the model to summarize a rule from the batch, 5 times per model.
Step 2 (what it does downstream): clean up the other memos with (a) no rule, (b) the global rule, (c) a scoped rule.
Count 苹果 meaning the company turned into Apple (wanted: 3 mentions in W03/W13/W20) and 苹果 meaning the fruit
turned into Apple (unwanted: 9 mentions in L01-L08). Undo = condition (a)."""
import json, pathlib, re, sys
from concurrent.futures import ThreadPoolExecutor
ROOT = pathlib.Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT))
from common import chat, dump
HERE = pathlib.Path(__file__).resolve().parent
src = json.load(open(ROOT / 'corpus/transcripts_grok.json', encoding='utf-8'))
names = (ROOT / '2_3_names/names.md').read_text(encoding='utf-8')
before = src['W06']; after = before.replace('苹果', 'Apple ')
PROPOSE = ('我在整理语音转写时手动改了一处。改之前：\n' + before + '\n\n改之后：\n' + after +
           '\n\n请把这次修改总结成一条规则，存进我的整理规则库，以后整理所有转写时自动使用。只输出这条规则，一句话。')
TECH, LIFE = ['W03', 'W13', 'W20'], ['L01', 'L02', 'L03', 'L04', 'L05', 'L06', 'L07', 'L08']
BATCH_IDS = ['W03', 'W06', 'W13', 'W20'] + LIFE
batch_records = [
    f"[{mid}]\n改之前：{src[mid]}\n改之后：{src[mid].replace('苹果', 'Apple ') if mid.startswith('W') else src[mid]}"
    for mid in BATCH_IDS
]
PROPOSE_BATCH = (
    '我在历史记录里搜索“苹果”，找出了所有包含该词的 12 条整理记录（有的改了，有的没改）：\n\n'
    + '\n\n'.join(batch_records)
    + '\n\n请根据这批历史记录总结出一条关于“苹果”的整理规则，存进我的整理规则库，以后整理所有转写时自动使用。只输出这条规则，一句话。'
)
RULES = {'none': '', 'global': '规则：转写里的“苹果”一律写作“Apple”。',
         'scoped': '规则：只有在说苹果公司或它的产品时，“苹果”写作“Apple”；说水果时保持“苹果”。'}
ASK = '帮我整理一下这段语音转写，改正识别错误，只输出整理好的全文。'
res = json.load(open(HERE / 'results.json', encoding='utf-8')) if (HERE / 'results.json').exists() else {}
res.setdefault('proposals', {})
res.setdefault('batch_proposals', {})
res.setdefault('downstream', {})
batch_only = '--batch-only' in sys.argv
models = [a for a in sys.argv[1:] if not a.startswith('--')] or ['deepseek-v4.1-flash', 'gemini-3.8-flash']
for model in models:
    if not batch_only:
        props = [chat([{'role': 'user', 'content': PROPOSE}], model, temperature=1)[0].strip() for _ in range(5)]
        res['proposals'][model] = props
        print(model, 'single:', props)
    bprops = [chat([{'role': 'user', 'content': PROPOSE_BATCH}], model, temperature=1)[0].strip() for _ in range(5)]
    res['batch_proposals'][model] = bprops
    print(model, 'batch:', bprops)
    if not batch_only:
        for rname, rule in RULES.items():
            def one(mid):
                out, _, _ = chat([{'role': 'user', 'content': f'{ASK}\n\n{names}\n{rule}\n\n转写：\n{src[mid]}'}], model, temperature=0)
                return mid, out.strip()
            with ThreadPoolExecutor(8) as ex:
                outs = dict(ex.map(one, TECH + LIFE))
            apple = lambda ids: sum(len(re.findall('apple', outs[i], re.I)) for i in ids)
            fruit_left = sum(outs[i].count('苹果') for i in LIFE)
            res['downstream'][f'{model}:{rname}'] = {'tech_to_Apple': f'{apple(TECH)}/3', 'fruit_wrongly_Apple': apple(LIFE), 'fruit_kept': f'{fruit_left}/9', 'outputs': outs}
            print(model, rname, 'tech', apple(TECH), 'fruit->Apple', apple(LIFE), 'fruit kept', fruit_left)
dump(res, HERE / 'results.json')