"""门禁（方案 §7.2 第 6 步）：新模型须 ①门禁集到达率不降 ②A 类降幅 ≤ 0.5 个百分点 ③上一轮修好的错题 95% 仍对。
  python -m ml.mining.gate --old cur.pt --new cand.pt --tasks ml/runs/gate_tasks.jsonl --out gate.json [--fixed prev_fixed.jsonl]
通过则把 cand 复制为 --current（默认 ml/runs/current.pt）。gate.json 含 accept / gain_pp，loop.sh 据此决定继续或停止。"""
import argparse, json, shutil
import torch

from ml.common.kin_torch import Kin
from ml.model.policy import build
from ml.sim.closed_loop import evaluate


def load(path, dev):
    ck = torch.load(path, map_location=dev)
    m = build(ck['cfg']).to(dev); m.load_state_dict(ck['model']); return m.eval()


def per_task_success(pers):
    return torch.cat([p['done'] & (p['col'] == 0) for _, p in pers]).tolist()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--old', required=True); ap.add_argument('--new', required=True); ap.add_argument('--tasks', required=True)
    ap.add_argument('--fixed', default=None); ap.add_argument('--out', required=True); ap.add_argument('--current', default='ml/runs/current.pt')
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    a = ap.parse_args()
    kin = Kin(device=a.device)
    tasks = [json.loads(l) for l in open(a.tasks) if l.strip()]
    ro, po = evaluate(load(a.old, a.device), kin, tasks, a.device)
    rn, pn = evaluate(load(a.new, a.device), kin, tasks, a.device)
    so, sn = per_task_success(po), per_task_success(pn)
    acc = lambda s, c: (sum(x for x, t in zip(s, tasks) if t['cat'] == c) / max(1, sum(t['cat'] == c for t in tasks)))
    gain = (sum(sn) - sum(so)) / len(tasks) * 100
    a_drop = (acc(so, 'A') - acc(sn, 'A')) * 100
    keep = None
    if a.fixed:
        ft = [json.loads(l) for l in open(a.fixed) if l.strip()]
        if ft:
            r2, p2 = evaluate(load(a.new, a.device), kin, ft, a.device)
            keep = sum(per_task_success(p2)) / len(ft)
    accept = gain >= 0 and a_drop <= 0.5 and (keep is None or keep >= 0.95)
    out = {'old': ro, 'new': rn, 'gain_pp': gain, 'A_drop_pp': a_drop, 'fixed_kept': keep, 'accept': accept}
    # 本轮新修好的题（旧错新对）留给下一轮做"仍对"检查
    fixed_now = [t for t, x, y in zip(tasks, so, sn) if (not x) and y]
    with open(a.out.replace('.json', '_fixed.jsonl'), 'w') as f:
        for t in fixed_now: f.write(json.dumps(t) + '\n')
    json.dump(out, open(a.out, 'w'), indent=1)
    if accept: shutil.copy(a.new, a.current)
    print(json.dumps({k: out[k] for k in ('gain_pp', 'A_drop_pp', 'fixed_kept', 'accept')}))


if __name__ == '__main__':
    main()