"""小规模闭环评估（部署口径）：
  python -m ml.eval.eval_cl --n 256 --commit 3 --fb 1.8 --partial [--ki 0.1 --select consist]
输出：到达 / 碰撞 / ratio2s / 刹停 / 精修接管 / 墙钟。直接调 rollout，参数与部署一致。"""
import argparse, functools, json, time
import torch

from ml.common.ckpt import load_policy
from ml.common.kin_torch import Kin
from ml.data import load_tasks
from ml.sim.closed_loop import rollout
import ml.sim.closed_loop as clmod


def run(model, kin, tasks, device, bs=256, **kw):
    agg, per = {}, []
    for i in range(0, len(tasks), bs):
        r, p = rollout(model, kin, tasks[i:i + bs], device, **kw)
        per.append((r, p, len(tasks[i:i + bs])))
    n = sum(w for _, _, w in per)
    for r, _, w in per:
        for k, v in r.items():
            if k in ('n', 'time_budget') or v is None: continue
            agg[k] = agg.get(k, 0.0) + v * w
    out = {k: v / n for k, v in agg.items()}; out['n'] = n
    out['goal_ratio_2s'] = torch.cat([p['goal_ratio_2s'] for _, p, _ in per]).median().item()
    out['err_fail_med'] = torch.cat([p['err'] for _, p, _ in per])[~torch.cat([p['ok'] for _, p, _ in per])].median().item()
    out['dls_share'] = torch.cat([p['dls_it'] for _, p, _ in per]).ge(0).float().mean().item()
    out['time_budget'] = per[0][0]['time_budget']
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ckpt', default='ml/runs/base_dag2_dag1/best.pt')
    ap.add_argument('--samples', default='ml/data/samples')
    ap.add_argument('--tasks', default=None)
    ap.add_argument('--split', default='val')
    ap.add_argument('--n', type=int, default=256)
    ap.add_argument('--K', type=int, default=8)
    ap.add_argument('--commit', type=int, default=3)
    ap.add_argument('--fb', type=float, default=1.8)
    ap.add_argument('--ki', type=float, default=0.0)
    ap.add_argument('--fb_scale', type=float, default=0.0)
    ap.add_argument('--fb_near', type=float, default=0.2)
    ap.add_argument('--fb_fallback', type=float, default=0.0)
    ap.add_argument('--kdecay', type=float, default=0.9)
    ap.add_argument('--kimax', type=float, default=0.5)
    ap.add_argument('--fb_every', type=int, default=1)
    ap.add_argument('--safety_try', type=int, default=0)
    ap.add_argument('--sel_w', type=float, default=0.1)
    ap.add_argument('--select', default='consist')
    ap.add_argument('--replan', default='r3')
    ap.add_argument('--safety', type=int, default=1, help='1=安全层+部分前缀；0=关')
    ap.add_argument('--partial', type=int, default=1)
    ap.add_argument('--max_sec', type=float, default=None)
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--dls', default='analytic', choices=['analytic', 'fd'])
    ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu')
    a = ap.parse_args()
    if a.dls == 'fd':
        clmod.dls_steps = functools.partial(clmod.dls_steps, analytic=False)
    model, ck = load_policy(a.ckpt, a.device)
    kin = Kin(device=a.device)
    assert ck.get('fp') in (None, kin.fp)
    tasks = ([json.loads(l) for l in open(a.tasks, encoding='utf-8') if l.strip()] if a.tasks
             else [t for t in load_tasks([a.samples]) if t['split'] == a.split])
    tasks = tasks[:a.n]
    t0 = time.perf_counter()
    r = run(model, kin, tasks, a.device, K=a.K, safety=bool(a.safety), seed=a.seed, max_sec=a.max_sec,
            safety_partial=(a.commit if (a.safety and a.partial) else 0), commit_n=a.commit,
            select=a.select, replan=a.replan, fb_gain=a.fb, fb_scale=a.fb_scale, fb_near=a.fb_near, fb_fallback=a.fb_fallback, fb_ki=a.ki, fb_decay=a.kdecay, fb_imax=a.kimax,
            fb_every=a.fb_every, safety_try=a.safety_try, sel_w=a.sel_w)
    wall = time.perf_counter() - t0
    print('ckpt=%s n=%d K=%d commit=%d fb=%.2f ki=%.2f select=%s replan=%s safety=%d' %
          (a.ckpt, r['n'], a.K, a.commit, a.fb, a.ki, a.select, a.replan, a.safety))
    print('success %.3f  time_ok_1p3 %.3f  coll %.3f  ratio2s %.3f  brakes %.1f  dls_share %.3f  err_fail %.2f  wall %.1fs (%.0f ms/题)'
          % (r['success'], r['time_ok_1p3'] or 0, r['collision_rate'], r['goal_ratio_2s'], r['brakes_mean'],
             r['dls_share'], r['err_fail_med'], wall, wall / max(1, len(tasks)) * 1000))


if __name__ == '__main__':
    main()