"""차량 객체 가우시안의 자세 합성 (nusim.recon.train._compose_objects, recon.objects) — CPU."""

import numpy as np
import pytest

torch = pytest.importorskip("torch")

from nusim.data.rig import matrix_to_quat, quat_to_matrix  # noqa: E402
from nusim.recon.train import _compose_objects  # noqa: E402


def _rand_q(rng, n):
    q = rng.normal(size=(n, 4))
    return q / np.linalg.norm(q, axis=1, keepdims=True)


def test_compose_objects_moves_only_object_gaussians():
    rng = np.random.default_rng(0)
    n, K = 50, 3
    means = rng.normal(size=(n, 3)).astype(np.float32)
    quats = _rand_q(rng, n).astype(np.float32)
    obj = np.array([-1] * 20 + [0] * 10 + [1] * 10 + [2] * 10, np.float32)
    params = {"means": torch.tensor(means), "quats": torch.tensor(quats),
              "opacities": torch.zeros(n), "obj": torch.tensor(obj)}
    Qb = _rand_q(rng, K)
    R = np.stack([quat_to_matrix(q) for q in Qb]).astype(np.float32)
    T = rng.normal(size=(K, 3)).astype(np.float32) * 10
    present = np.array([1, 1, 0], np.float32)  # 객체 2 는 이 시각에 박스가 없다
    m, q, op = _compose_objects(params, tuple(torch.tensor(a) for a in (R, T, Qb.astype(np.float32), present)))
    m, q, op = m.numpy(), q.numpy(), op.numpy()
    bg = obj < 0
    assert np.allclose(m[bg], means[bg]) and np.allclose(q[bg], quats[bg]) and np.allclose(op[bg], 0.5)
    for i in np.nonzero(~bg)[0]:
        k = int(obj[i])
        assert np.allclose(m[i], R[k] @ means[i] + T[k], atol=1e-5)
        # 회전 합성: 세계 회전 = 박스 회전 · 로컬 회전
        assert np.allclose(quat_to_matrix(q[i]), R[k] @ quat_to_matrix(quats[i]), atol=1e-5)
        assert op[i] == pytest.approx(0.5 * present[k])


def test_compose_objects_without_objects_is_identity():
    params = {"means": torch.zeros(4, 3), "quats": torch.tensor([[1.0, 0, 0, 0]] * 4), "opacities": torch.zeros(4)}
    m, q, op = _compose_objects(params, None)
    assert torch.equal(m, params["means"]) and torch.allclose(op, torch.full((4,), 0.5))


def test_matrix_quat_consistency_for_box_rotation():
    # objects 가 쓰는 규약 (quat_to_matrix) 이 rig.py 와 같은지 — 박스 회전 왕복
    q = np.array([0.9, 0.1, -0.2, 0.37])
    q /= np.linalg.norm(q)
    assert np.allclose(quat_to_matrix(matrix_to_quat(quat_to_matrix(q))), quat_to_matrix(q))