"""SE(3) and quaternion utilities.

Conventions
-----------
* Quaternions are **scalar-first** ``[w, x, y, z]`` to match gsplat's
  rasterization API (so gaussian rotations flow straight to the rasterizer).
* Camera coordinates are **OpenCV** style: x-right, y-down, z-forward; depth is
  the camera-space z. A ``cam2world`` (a.k.a. extrinsic / pose) matrix maps a
  point in camera coordinates to world coordinates; ``world2cam`` is its inverse
  and is what the rasterizer consumes as ``viewmats``.
"""

from __future__ import annotations

import torch


def normalize_quat(quat: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """Normalize a ``[..., 4]`` quaternion to unit norm."""
    return quat / quat.norm(dim=-1, keepdim=True).clamp_min(eps)


def quat_to_rotmat(quat: torch.Tensor) -> torch.Tensor:
    """Convert scalar-first quaternions ``[..., 4]`` to rotations ``[..., 3, 3]``."""
    quat = normalize_quat(quat)
    w, x, y, z = quat.unbind(-1)
    # Standard wxyz -> R.
    tx, ty, tz = 2 * x, 2 * y, 2 * z
    twx, twy, twz = tx * w, ty * w, tz * w
    txx, txy, txz = tx * x, ty * x, tz * x
    tyy, tyz, tzz = ty * y, tz * y, tz * z
    R = torch.stack(
        [
            1 - (tyy + tzz), txy - twz, txz + twy,
            txy + twz, 1 - (txx + tzz), tyz - twx,
            txz - twy, tyz + twx, 1 - (txx + tyy),
        ],
        dim=-1,
    )
    return R.reshape(quat.shape[:-1] + (3, 3))


def rotmat_to_quat(R: torch.Tensor) -> torch.Tensor:
    """Convert rotations ``[..., 3, 3]`` to scalar-first quaternions ``[..., 4]``."""
    m = R.reshape(R.shape[:-2] + (9,))
    m00, m01, m02, m10, m11, m12, m20, m21, m22 = m.unbind(-1)
    trace = m00 + m11 + m22
    # Branchless-ish via the four candidate formulations; pick the most stable.
    q0 = torch.stack([1 + trace, m21 - m12, m02 - m20, m10 - m01], dim=-1)
    q1 = torch.stack([m21 - m12, 1 + m00 - m11 - m22, m01 + m10, m02 + m20], dim=-1)
    q2 = torch.stack([m02 - m20, m01 + m10, 1 - m00 + m11 - m22, m12 + m21], dim=-1)
    q3 = torch.stack([m10 - m01, m02 + m20, m12 + m21, 1 - m00 - m11 + m22], dim=-1)
    cond = torch.stack([trace, m00, m11, m22], dim=-1)
    idx = cond.argmax(dim=-1, keepdim=True)
    stacked = torch.stack([q0, q1, q2, q3], dim=-2)  # [..., 4, 4]
    quat = torch.gather(stacked, -2, idx.unsqueeze(-1).expand(idx.shape + (4,))).squeeze(-2)
    return normalize_quat(quat)


def quat_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """Hamilton product of two scalar-first quaternions."""
    aw, ax, ay, az = a.unbind(-1)
    bw, bx, by, bz = b.unbind(-1)
    return torch.stack(
        [
            aw * bw - ax * bx - ay * by - az * bz,
            aw * bx + ax * bw + ay * bz - az * by,
            aw * by - ax * bz + ay * bw + az * bx,
            aw * bz + ax * by - ay * bx + az * bw,
        ],
        dim=-1,
    )


def se3_inverse(T: torch.Tensor) -> torch.Tensor:
    """Invert a ``[..., 4, 4]`` rigid transform."""
    R = T[..., :3, :3]
    t = T[..., :3, 3:]
    Rt = R.transpose(-1, -2)
    out = torch.zeros_like(T)
    out[..., :3, :3] = Rt
    out[..., :3, 3:] = -Rt @ t
    out[..., 3, 3] = 1.0
    return out


def apply_se3(T: torch.Tensor, points: torch.Tensor) -> torch.Tensor:
    """Apply ``[..., 4, 4]`` transform to ``[..., N, 3]`` points (broadcasting on batch)."""
    R = T[..., :3, :3]
    t = T[..., :3, 3]
    return torch.einsum("...ij,...nj->...ni", R, points) + t.unsqueeze(-2)