"""DLS 解析雅可比对拍 + 延迟（python -m ml.eval.check_dls）"""
import time
import torch

from ml.common.kin_torch import Kin
from ml.sim.closed_loop import dls_steps


def main():
    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
    kin = Kin(device=dev)
    g = torch.Generator(device=dev); g.manual_seed(0)
    B = 64
    lo, hi = kin.lo + 20, kin.hi - 20
    q = lo + (hi - lo) * torch.rand(B, 6, device=dev, generator=g)
    xf = kin.cum_xf(q)
    p = kin.fk(q, xf)
    T = p + 6 * torch.randn(B, 3, device=dev, generator=g)
    R = torch.zeros(B, 9, device=dev)
    a = dls_steps(kin, q, T, R, n=3, analytic=True)
    b = dls_steps(kin, q, T, R, n=3, analytic=False)
    d = (a - b).abs()
    print('pos-only  max|Δdq| %.6f°  mean %.6f°  rel %.4f' % (d.max().item(), d.mean().item(), (d.norm() / b.norm()).item()))
    # 带姿态目标：绕随机轴 ±8°
    axis = torch.randn(B, 3, device=dev, generator=g); axis = axis / axis.norm(dim=-1, keepdim=True)
    ang = (torch.rand(B, 1, device=dev, generator=g) * 16 - 8) * torch.pi / 180
    Kx = torch.zeros(B, 3, 3, device=dev)
    Kx[:, 0, 1] = -axis[:, 2]; Kx[:, 0, 2] = axis[:, 1]; Kx[:, 1, 0] = axis[:, 2]
    Kx[:, 1, 2] = -axis[:, 0]; Kx[:, 2, 0] = -axis[:, 1]; Kx[:, 2, 1] = axis[:, 0]
    I = torch.eye(3, device=dev)
    Rd = I + torch.sin(ang)[..., None] * Kx + (1 - torch.cos(ang))[..., None] * (Kx @ Kx)
    Rt = kin.orient(q, xf) @ Rd
    R9 = Rt.reshape(B, 9)
    a = dls_steps(kin, q, T, R9, n=3, analytic=True)
    b = dls_steps(kin, q, T, R9, n=3, analytic=False)
    d = (a - b).abs()
    print('pos+rot   max|Δdq| %.6f°  mean %.6f°  rel %.4f' % (d.max().item(), d.mean().item(), (d.norm() / b.norm()).item()))

    for n_ in (1, 3):
        for Bn in (1, 256):
            qq = q[:Bn]; TT = T[:Bn]; RR = R9[:Bn]
            for an in (False, True):
                for _ in range(3):
                    dls_steps(kin, qq, TT, RR, n=n_, analytic=an)
                torch.cuda.synchronize() if dev == 'cuda' else None
                t0 = time.perf_counter()
                for _ in range(20):
                    dls_steps(kin, qq, TT, RR, n=n_, analytic=an)
                torch.cuda.synchronize() if dev == 'cuda' else None
                print('n=%d B=%3d analytic=%d  %.2f ms' % (n_, Bn, an, (time.perf_counter() - t0) / 20 * 1e3))


if __name__ == '__main__':
    main()