"""小规模同步闭环评估；仿真墙钟不等于控制延迟。"""
import argparse, functools, hashlib, json, math, time
from pathlib import Path
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, tasks_to_tensors
import ml.sim.closed_loop as clmod


def run(model, kin, tasks, device, bs=256, **kw):
    per = []
    for i in range(0, len(tasks), bs):
        t0 = time.perf_counter()
        print('batch %d-%d/%d start' % (i + 1, min(i + bs, len(tasks)), len(tasks)), flush=True)
        def progress(step, total):
            print('batch %d step %d/%d wall %.1fs' % (i // bs + 1, step, total, time.perf_counter() - t0), flush=True)
        r, p = rollout(model, kin, tasks[i:i + bs], device, progress=progress, **kw)
        per.append((r, p, len(tasks[i:i + bs])))
        print('batch %d done wall %.1fs' % (i // bs + 1, time.perf_counter() - t0), flush=True)
    flat = {k: torch.cat([p[k] for _, p, _ in per]) for k in per[0][1] if k not in ('traj', 'plans')}
    _, _, _, ex, tr = tasks_to_tensors(tasks, 'cpu')
    ok, col = flat['ok'], flat['col'] > 0
    has_ref = torch.isfinite(ex)
    okr = ok & has_ref & torch.isfinite(tr)
    median = lambda x: x.median().item() if x.numel() else None
    out = {
        'n': len(tasks), 'n_ref': int(has_ref.sum()),
        'success': ok.float().mean().item(),
        'collision_free_success': (ok & ~col).float().mean().item(),
        'time_ok_1p3': (ok & has_ref & (flat['t_done'] <= 1.3 * ex)).sum().item() / int(has_ref.sum()) if has_ref.any() else None,
        'collision_rate': col.float().mean().item(),
        'collision_steps_mean': flat['col'].mean().item(),
        'travel_ratio_median': median(flat['travel'][okr] / tr[okr].clamp(min=1)),
        'time_mean_s': flat['t_done'][ok].mean().item() if ok.any() else None,
        'brakes_mean': flat['brakes'].mean().item(),
        'start_fail': flat['start_fail'].float().mean().item(),
        'goal_ratio_2s': median(flat['goal_ratio_2s']),
        'err_fail_med': median(flat['err'][~ok]),
        'err_success_med': median(flat['err'][ok]),
        'dls_share': flat['dls_it'].ge(0).float().mean().item(),
        'time_budget': ('fixed %.3f s' % kw['max_sec']) if kw.get('max_sec') is not None else
                       (per[0][0]['time_budget'] if len({r['time_budget'] for r, _, _ in per}) == 1 else 'mixed'),
    }
    if 'speed_ratio' in flat:   # 实际执行速度：每题时限内 max|Δq|/STEP（1.0 = VMAX），加速度同口径按每控制步的速度变化
        v, acc = flat['speed_ratio'], flat['accel_ratio']
        out.update({'speed_ratio_p50': median(v), 'speed_ratio_p95': v.quantile(0.95).item(),
                    'speed_ratio_max': v.max().item(), 'over_vmax_rate': (v > 1.001).float().mean().item(),
                    'accel_ratio_p95': acc.quantile(0.95).item(), 'accel_ratio_max': acc.max().item()})
    return out, per


def json_finite(value):
    if isinstance(value, dict): return {k: json_finite(v) for k, v in value.items()}
    if isinstance(value, list): return [json_finite(v) for v in value]
    if isinstance(value, float) and not math.isfinite(value): return None
    return value


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('--dls_safety', type=int, default=0)
    ap.add_argument('--fb_clear', type=float, default=0.0)
    ap.add_argument('--fb_cmin', type=float, default=0.3)
    ap.add_argument('--sub_refresh', type=int, default=0)
    ap.add_argument('--fb_block_gain', type=float, default=1.0)
    ap.add_argument('--dump', default=None)
    ap.add_argument('--out', default=None)
    ap.add_argument('--trace', default=None)
    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('--inc_cap', type=int, default=0, help='1=反馈叠加后再按每轴 VMAX·DT 限幅')
    ap.add_argument('--gap_margin', type=float, default=0.0, help='安全层要求的间隙余量 mm（>0 时宽相位 cap 自动抬到 margin+0.5）')
    ap.add_argument('--cert_batch', type=int, default=0, help='1=一次 gaps 调用覆盖全部候选（压最坏情况，典型路径略贵）')
    ap.add_argument('--steps', type=int, default=1, help='MeanFlow 采样段数；1=线上一口径（历史发布数字全是 1）')
    ap.add_argument('--backoff', type=int, default=0, help='全拒时把首选候选按 1/2、1/4…缩步重试的档数（需 --cert_batch 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()
    print('config ' + json.dumps(vars(a), ensure_ascii=False), flush=True)
    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]
    if not tasks: ap.error('no tasks selected')
    meta = {
        'config': vars(a), 'batch_size': 256, 'batch_seed_rule': 'same seed per batch',
        'checkpoint_sha256': hashlib.sha256(Path(a.ckpt).read_bytes()).hexdigest(),
        'tasks_sha256': hashlib.sha256(json.dumps(tasks, sort_keys=True).encode()).hexdigest(),
        'kin_fp': kin.fp, 'act_scale': float(model.act_scale),
    }
    t0 = time.perf_counter()
    r, pers = run(model, kin, tasks, a.device, K=a.K, safety=bool(a.safety), seed=a.seed, max_sec=a.max_sec, steps=a.steps,
            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, dls_safety=a.dls_safety, fb_clear=a.fb_clear, fb_cmin=a.fb_cmin, sub_refresh=a.sub_refresh, fb_block_gain=a.fb_block_gain,
            inc_cap=a.inc_cap, gap_margin=a.gap_margin, cert_batch=a.cert_batch, backoff=a.backoff, return_traj=bool(a.trace))
    wall = time.perf_counter() - t0
    result = {**meta, 'metrics': r, 'simulation_wall_s': wall}
    print(json.dumps(json_finite(result), ensure_ascii=False, allow_nan=False), flush=True)
    if a.out:
        Path(a.out).write_text(json.dumps(json_finite(result), indent=2, ensure_ascii=False, allow_nan=False) + '\n', encoding='utf-8')
    if a.dump:
        rows = []
        keys = ['done', 'ok', 't_done', 'limit', 'col', 'brakes', 'err', 'first_hit', 'dls_it', 'goal_ratio_2s']
        flat = {k: torch.cat([p[k] for _, p, _ in pers]).tolist() for k in keys}
        for i, t in enumerate(tasks):
            rows.append({**t, 'i': i, **{k: flat[k][i] for k in keys}})
        Path(a.dump).write_text(json.dumps(json_finite(rows), allow_nan=False), encoding='utf-8')
        print('dump', a.dump, flush=True)
    if a.trace:
        torch.save({**meta, 'tasks': tasks, 'batches': pers}, a.trace)
        print('trace', a.trace, flush=True)


if __name__ == '__main__':
    main()