"""对拍：python -m ml.eval.check_parity [--kin ml/assets/kin.json] [--data ml/data/parity.json]
检查 PyTorch 版 fk / orient / 组对间隙 / 83 维特征 与 JS 一致，以及间隙梯度与有限差分一致。"""
import argparse, json
import torch
from ml.common.kin_torch import Kin
from ml.common.features import features, GAP_CAP


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--kin', default=None); ap.add_argument('--data', default='ml/data/parity.json')
    a = ap.parse_args()
    kin = Kin(a.kin) if a.kin else Kin(dtype=torch.float64)
    D = json.load(open(a.data, encoding='utf-8'))
    assert D['fp'] == kin.fp, f"指纹不符：parity {D['fp']} vs kin {kin.fp}（重跑 export_kin.cjs）"
    R = D['recs']; f64 = torch.float64
    q = torch.tensor([r['q'] for r in R], dtype=f64)
    qp = torch.tensor([r['qPrev'] for r in R], dtype=f64)
    T = torch.tensor([r['T'] for r in R], dtype=f64)
    Rt = torch.tensor([r['R'] if r['R'] else [0.] * 9 for r in R], dtype=f64)
    pre = torch.zeros(len(R), 3, 6, dtype=f64); pl = torch.zeros(len(R), dtype=torch.long)
    for n, r in enumerate(R):
        for k, d in enumerate(r['prefix']): pre[n, k] = torch.tensor(d, dtype=f64)
        pl[n] = len(r['prefix'])
    sub = torch.tensor([r['sub'] if r['sub'] else [0.] * 6 for r in R], dtype=f64)
    sf = torch.tensor([1. if r['sub'] else 0. for r in R], dtype=f64)
    kin = Kin(a.kin, dtype=f64) if a.kin else kin
    xf = kin.cum_xf(q)
    e_fk = (kin.fk(q, xf) - torch.tensor([r['fk'] for r in R], dtype=f64)).abs().max().item()
    e_or = (kin.orient(q, xf).reshape(-1, 9) - torch.tensor([r['orient'] for r in R], dtype=f64)).abs().max().item()
    G = kin.gaps(q, GAP_CAP, xf)
    gm_js = torch.tensor([r['gm'] for r in R], dtype=f64)
    e_gm = (G['gm'] - gm_js.clamp(max=GAP_CAP)).abs().max().item()
    col_js = gm_js.min(-1).values <= 0; col_py = G['gm'].min(-1).values <= 0
    f = features(kin, q, qp, T, Rt, pre, pl, sub, sf, gm=G['gm'], xf=xf)
    e_f = (f - torch.tensor([r['feat'] for r in R], dtype=f64)).abs().max().item()
    # 梯度：最小间隙对 q 的梯度 vs 有限差分（取 50 个无干涉且间隙 < 20 的样本）
    sel = ((G['gm'].min(-1).values > 0.5) & (G['gm'].min(-1).values < 20)).nonzero().flatten()[:50]
    rel = []
    if len(sel):
        qs = q[sel].clone().requires_grad_(True)
        g = kin.gaps(qs, GAP_CAP)['gm'].min(-1).values
        gr, = torch.autograd.grad(g.sum(), qs)
        h = 1e-4
        for j in range(6):
            dq = torch.zeros_like(qs); dq[:, j] = h
            fd = (kin.gaps(qs.detach() + dq, GAP_CAP)['gm'].min(-1).values - kin.gaps(qs.detach() - dq, GAP_CAP)['gm'].min(-1).values) / (2 * h)
            rel.append(((gr[:, j] - fd).abs() / (fd.abs() + 1e-3)).max().item())
    res = {'n': len(R), 'fk_max_abs_mm': e_fk, 'orient_max_abs': e_or, 'gap_max_abs_mm': e_gm,
           'collide_bool_mismatch': int((col_js != col_py).sum()), 'feat_max_abs': e_f,
           'grad_rel_err_max': max(rel) if rel else None}
    ok = e_fk < 1e-6 and e_or < 1e-9 and e_gm < 1e-4 and res['collide_bool_mismatch'] == 0 and e_f < 1e-4 and (not rel or max(rel) < 1e-3)
    res['PASS'] = bool(ok)
    print(json.dumps(res, indent=1))


if __name__ == '__main__':
    main()