import torch
import numpy as np
import cv2

from vidar.geometry.cameras.base import CameraBase


class CameraGoPro(CameraBase):
    def __init__(self, K, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.geometry = 'gopro'

        self._K = K
        D = torch.tensor(
            [0.036902, 0.057131, -0.057154, 0.018458],
            dtype=K.dtype, device=K.device
        )
        self.D = D.unsqueeze(0).repeat(K.shape[0], 1)
        

    def pixelToRay(self, grid):

        rays = []
        for b in range(grid.shape[0]):
            uv = grid[b, :2, :]
            uv = uv.unsqueeze(1).permute(2, 1, 0)
            undistorted = cv2.fisheye.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(rays, dtype=grid.dtype, device=grid.device)
        rays = torch.cat([rays, torch.ones_like(rays[:, :1, :])], dim=1)

        return rays


    def lift(self, grid=None, depth=None):
        """
        Given a grid of uv coordinates and the associated depth, get the 3D points in
        the camera's space.
        grid: [B, 3, HW]
        depth: [B, 1, HW]
        """

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

        xyz = rays * depth
    
        return xyz.float()