import math

import torch

from mapgs.geometry import (
    quat_to_rotmat, rotmat_to_quat, se3_inverse, project_points,
    backproject_depth, plucker_embedding, look_at_pose,
)


def test_quat_rotmat_roundtrip():
    q = torch.randn(20, 4)
    q = q / q.norm(dim=-1, keepdim=True)
    R = quat_to_rotmat(q)
    # orthonormal, det 1
    eye = R @ R.transpose(-1, -2)
    assert torch.allclose(eye, torch.eye(3).expand_as(eye), atol=1e-5)
    assert torch.allclose(torch.det(R), torch.ones(20), atol=1e-5)
    q2 = rotmat_to_quat(R)
    R2 = quat_to_rotmat(q2)
    assert torch.allclose(R, R2, atol=1e-5)


def test_se3_inverse():
    T = torch.eye(4).repeat(5, 1, 1)
    T[:, :3, :3] = quat_to_rotmat(torch.randn(5, 4))
    T[:, :3, 3] = torch.randn(5, 3)
    I = T @ se3_inverse(T)
    assert torch.allclose(I, torch.eye(4).expand_as(I), atol=1e-5)


def test_projection_roundtrip():
    K = torch.tensor([[300., 0, 224], [0, 300, 128], [0, 0, 1]])[None]
    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 20, 0.5]))[None]
    pts = torch.randn(1, 100, 3)
    pts[..., 1] = pts[..., 1].abs() + 5  # in front
    uv, z = project_points(pts, K, se3_inverse(c2w))
    back = backproject_depth(uv, z, K, c2w)
    assert torch.allclose(back, pts, atol=1e-3)
    assert (z > 0).all()


def test_look_at_forward_is_positive_depth():
    # a point in front of a forward-looking camera must have positive cam-z
    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 30, 0.5]))[None]
    K = torch.tensor([[300., 0, 224], [0, 300, 128], [0, 0, 1]])[None]
    pt = torch.tensor([[[0., 10, 0.5]]])
    _, z = project_points(pt, K, se3_inverse(c2w))
    assert z.item() > 0


def test_plucker_shape():
    K = torch.tensor([[300., 0, 224], [0, 300, 128], [0, 0, 1]])[None].repeat(3, 1, 1)
    c2w = torch.eye(4)[None].repeat(3, 1, 1)
    pl = plucker_embedding(K, c2w, 32, 48)
    assert pl.shape == (3, 6, 32, 48)