import torch
import numpy as np
import cv2

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


class CameraDistorted(CameraBase):
    """Distorted camera using OpenCV fisheye or plumb_bob undistortion.
    K: [B,8+] as (fx, fy, cx, cy, d0, ...) or [B,3,3] with separate D.
    Geometry auto-detected from distortion param count (4=fisheye, 5/12=plumb_bob)."""

    def __init__(self, K, D=None, *args, **kwargs):
        super().__init__(*args, **kwargs)

        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
            self._D = K[:, 4:]
        else:
            # It's ready
            self._K = K
            self._D = D

        if self._D.shape[1] == 4:
            # fx, fy, cx, cy, k1, k2, k3, k4
            self.geometry = 'fisheye'
        elif self._D.shape[1] == 5:
            # fx, fy, cx, cy, k1, k2, p1, p2, k3
            self.geometry = 'plumb_bob'
        elif self._D.shape[1] == 12:
            # fx, fy, cx, cy, k1, k2, p1, p2, k3, k4, k5, k6, s1, s2, s3, s4
            self.geometry = 'plumb_bob'
        else: 
            raise ValueError(f'Invalid distortion parameters {self._D}')

    # 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 OpenCV.
        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)
            if self.geometry == 'fisheye':
                undistorted = cv2.fisheye.undistortPoints(
                    uv.cpu().numpy(),
                    self._K[b].cpu().numpy(),
                    self._D[b].cpu().numpy(),
                    P=np.eye(3),
                )
            elif self.geometry == 'plumb_bob':
                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 OpenCV.
        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()
            if self.geometry == 'fisheye':
                projected_b, _ = cv2.fisheye.projectPoints(
                    objectPoints=p,
                    rvec=np.zeros(3), 
                    tvec=np.zeros(3),
                    K=self._K[b].cpu().numpy(),
                    D=self._D[b].cpu().numpy(),
                )
            elif self.geometry == 'plumb_bob':
                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))
        depth[points[:, 2] < 0] = 0
        coords = self.rayToPixel(points)

        return coords, depth