import torch

from anydata.geometry.cameras.base import CameraBase
from anydata.utils.data import cat_channel_ones




class CameraPinhole(CameraBase):
    """Standard pinhole camera model.
    K: [B,3,3] intrinsics or [B,4] as (fx, fy, cx, cy)."""

    def __init__(self, K, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.geometry = 'pinhole'

        if len(K.shape) == 2 and K.shape[1] == 4:
            # Parse intrinsics from Bx4 to Bx3x3
            self._K = torch.eye(3, dtype=K.dtype, device=K.device).unsqueeze(0).repeat(K.shape[0], 1, 1)
            self._K[:, 0, 0] = K[:, 0]
            self._K[:, 1, 1] = K[:, 1]
            self._K[:, 0, 2] = K[:, 2]
            self._K[:, 1, 2] = K[:, 3]
        elif len(K.shape) == 3 and K.shape[1] == 3 and K.shape[2] == 3:
            # Intrinsics is already Bx3x3
            self._K = K


    def Pwc(self, from_world=True):
        """Projection matrix [B,4,4]: K (padded to 4x4) @ Twc.
        If from_world=False, returns just K padded to 4x4."""
        K = torch.eye(4, device=self.device, dtype=self.dtype).unsqueeze(0).repeat(self.b, 1, 1)
        K[:, :3, :3] = self._K
        if not from_world or self.Twc is None:
            return K
        return torch.bmm(K, self.Twc)

    def lift(self, grid, depth):
        """Lift 2D pixel coords + depth to 3D camera-frame points.
        grid: [B,3,N] (pixel coords with ones), depth: [B,1,N]. Returns [B,3,N]."""
        rays = grid.clone()

        # NOTE(bvh): former bug, unsqueeze(-1) needed for correct broadcasting when B > 1
        rays[:, 0] = (rays[:, 0] - self.K[:, 0, 2].unsqueeze(-1)) / self.K[:, 0, 0].unsqueeze(-1)
        rays[:, 1] = (rays[:, 1] - self.K[:, 1, 2].unsqueeze(-1)) / self.K[:, 1, 1].unsqueeze(-1)

        return rays * depth

    def unlift(self, points, from_world=True, euclidean=False):
        """Project 3D points to 2D pixel coords + depth.
        points: [B,3,N]. Returns (coords [B,2,N], depth [B,N])."""
        projected = torch.bmm(self.Pwc(from_world)[:, :3], cat_channel_ones(points, 1))
        coords = projected[:, :2] / torch.clamp(projected[:, 2].unsqueeze(1), min=1e-7)
        if not euclidean:
            depth = projected[:, 2]
        else:
            points = self.from_world(points) if from_world else points
            depth = torch.linalg.vector_norm(points, dim=1, keepdim=False)
        depth[projected[:, 2] <= 0] = 0
        return coords, depth