import torch
import numpy as np
import cv2

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


class CameraPlumbBob(CameraBase):
    """Plumb-bob distortion model using cv2.undistortPoints / cv2.projectPoints.
    K: [B,9] as (fx, fy, cx, cy, k1, k2, p1, p2, k3) or [B,3,3] with separate D."""

    def __init__(self, K, D=None, *args, **kwargs):
        self.geometry = 'plumb_bob'

        if D is None:
            # Parse intrinsics
            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]
            # Parse distortion (5 params: k1, k2, p1, p2, k3)
            self._D = K[:, 4:]
        else:
            # It's ready
            self._K = K
            self._D = D

        super().__init__(*args, **kwargs)

    # NOTE(bvh): override CameraBase methods since we must also pass D

    def __getitem__(self, idx):
        if isinstance(idx, int):
            idx = slice(idx, idx + 1)
        return type(self)(K=self._K[idx], D=self._D[idx], hw=self._hw, Tcw=self._Tcw[idx])

    def clone(self):
        return type(self)(K=self._K.clone(), D=self._D.clone(), hw=self._hw,
                          Tcw=self._Tcw.clone() if self._Tcw is not None else None)

    @staticmethod
    def from_list(cams):
        K = torch.cat([c._K for c in cams], 0)
        D = torch.cat([c._D for c in cams], 0)
        Tcw = torch.cat([c.Tcw for c in cams], 0)
        return type(cams[0])(K=K, D=D, hw=cams[0].hw, Tcw=Tcw)

    def no_translation(self):
        Tcw = self.Tcw.clone()
        Tcw[:, :-1, -1] = 0
        return type(self)(K=self._K, D=self._D, hw=self._hw, Tcw=Tcw)

    def pixelToRay(self, grid):
        """Convert pixel coordinates to undistorted ray directions using cv2.undistortPoints.
        grid: [B,2or3,N]. Returns [B,3,N]."""

        rays = []
        for b in range(grid.shape[0]):
            uv = grid[b, :2, :]
            uv = uv.unsqueeze(1).permute(2, 1, 0)
            undistorted = cv2.undistortPoints(
                uv.cpu().numpy(),
                self._K[b].cpu().numpy(),
                self._D[b].cpu().numpy(),
                P=np.eye(3),
            )
            undistorted = undistorted.squeeze(1).transpose()
            rays.append(undistorted)

        rays = torch.tensor(np.array(rays), dtype=grid.dtype, device=grid.device)
        rays = torch.cat([rays, torch.ones_like(rays[:, :1, :])], dim=1)

        return rays


    def rayToPixel(self, points):
        """Project 3D camera-frame points to distorted pixel coordinates using cv2.projectPoints.
        points: [B,3,N]. Returns [B,2,N]."""

        projected = []
        for b in range(points.shape[0]):
            p = points[b].permute(1, 0).unsqueeze(1).cpu().numpy()
            projected_b, _ = cv2.projectPoints(
                objectPoints=p,
                rvec=np.zeros(3),
                tvec=np.zeros(3),
                cameraMatrix=self._K[b].cpu().numpy(),
                distCoeffs=self._D[b].cpu().numpy(),
            )
            projected.append(torch.tensor(projected_b).squeeze(1).permute(1, 0))

        projected = torch.stack(projected, 0)
        projected = projected.to(device=points.device, dtype=points.dtype)

        return projected


    def lift(self, grid=None, depth=None):
        """Lift 2D pixel coords + depth to 3D using undistorted rays.
        grid: [B,3,N], depth: [B,1,N]. Returns [B,3,N]."""

        rays = self.pixelToRay(grid)
        rays[torch.isnan(rays)] = 0.0
        rays = torch.nn.functional.normalize(rays, dim=1, p=2)

        return (rays * depth).float()


    def unlift(self, points, from_world=True, euclidean=False):
        """Project 3D points to distorted 2D pixel coords + euclidean depth.
        points: [B,3,N]. Returns (coords [B,2,N], depth [B,N])."""
        points = torch.matmul(self.Twc[:, :3], cat_channel_ones(points, 1))
        depth = torch.sqrt((points ** 2).sum(1))
        coords = self.rayToPixel(points)
        return coords, depth