#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the

# LICENSE file in the root directory of this source tree.

r"""This module defines various ObservationTransformers that can be used
to transform the output of the simulator before they are fed into the
policy of the neural network. This can include various useful preprocessing
including faking a semantic sensor using RGB input and MaskRCNN or faking
a depth sensor using RGB input. You can also stich together multiple sensors.
This code runs on the batched of inputs to these networks efficiently.
ObservationTransformer all run as nn.modules and can be used for encoders or
any other neural networks preprocessing steps.
Assumes the input is on CUDA.

They also implement a function that transforms that observation space so help
fake or modify sensor input from the simulator.

This module API is experimental and likely to change
"""
import abc
import copy
import numbers
from enum import Enum
from typing import Dict, Iterable, List, Optional, Tuple, Union

import numpy as np
import torch
from gym import spaces
from torch import nn

from habitat.config import Config
from habitat.core.logging import logger
from habitat_baselines.common.baseline_registry import baseline_registry
from habitat_baselines.utils.common import (
    center_crop,
    get_image_height_width,
    image_resize_shortest_edge,
    overwrite_gym_box_shape,
)


class ObservationTransformer(nn.Module, metaclass=abc.ABCMeta):
    """This is the base ObservationTransformer class that all other observation
    Transformers should extend. from_config must be implemented by the transformer.
    transform_observation_space is only needed if the observation_space ie.
    (resolution, range, or num of channels change)."""

    def transform_observation_space(
        self, observation_space: spaces.Dict, **kwargs
    ):
        return observation_space

    @classmethod
    @abc.abstractmethod
    def from_config(cls, config: Config):
        pass

    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        return observations


@baseline_registry.register_obs_transformer()
class ResizeShortestEdge(ObservationTransformer):
    r"""An nn module the resizes your the shortest edge of the input while maintaining aspect ratio.
    This module assumes that all images in the batch are of the same size.
    """

    def __init__(
        self,
        size: int,
        channels_last: bool = True,
        trans_keys: Tuple[str] = ("rgb", "depth", "semantic"),
    ):
        """Args:
        size: The size you want to resize the shortest edge to
        channels_last: indicates if channels is the last dimension
        """
        super(ResizeShortestEdge, self).__init__()
        self._size: int = size
        self.channels_last: bool = channels_last
        self.trans_keys: Tuple[str] = trans_keys

    def transform_observation_space(
        self,
        observation_space: spaces.Dict,
    ):
        size = self._size
        observation_space = copy.deepcopy(observation_space)
        if size:
            for key in observation_space.spaces:
                if key in self.trans_keys:
                    # In the observation space dict, the channels are always last
                    h, w = get_image_height_width(
                        observation_space.spaces[key], channels_last=True
                    )
                    if size == min(h, w):
                        continue
                    scale = size / min(h, w)
                    new_h = int(h * scale)
                    new_w = int(w * scale)
                    new_size = (new_h, new_w)
                    logger.info(
                        "Resizing observation of %s: from %s to %s"
                        % (key, (h, w), new_size)
                    )
                    observation_space.spaces[key] = overwrite_gym_box_shape(
                        observation_space.spaces[key], new_size
                    )
        return observation_space

    def _transform_obs(self, obs: torch.Tensor) -> torch.Tensor:
        return image_resize_shortest_edge(
            obs, self._size, channels_last=self.channels_last
        )

    @torch.no_grad()
    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        if self._size is not None:
            observations.update(
                {
                    sensor: self._transform_obs(observations[sensor])
                    for sensor in self.trans_keys
                    if sensor in observations
                }
            )
        return observations

    @classmethod
    def from_config(cls, config: Config):
        return cls(config.RL.POLICY.OBS_TRANSFORMS.RESIZE_SHORTEST_EDGE.SIZE)


@baseline_registry.register_obs_transformer()
class CenterCropper(ObservationTransformer):
    """An observation transformer is a simple nn module that center crops your input."""

    def __init__(
        self,
        size: Union[int, Tuple[int, int]],
        channels_last: bool = True,
        trans_keys: Tuple[str] = ("rgb", "depth", "semantic"),
    ):
        """Args:
        size: A sequence (h, w) or int of the size you wish to resize/center_crop.
                If int, assumes square crop
        channels_list: indicates if channels is the last dimension
        trans_keys: The list of sensors it will try to centercrop.
        """
        super().__init__()
        if isinstance(size, numbers.Number):
            size = (int(size), int(size))
        assert len(size) == 2, "forced input size must be len of 2 (h, w)"
        self._size = size
        self.channels_last = channels_last
        self.trans_keys = trans_keys  # TODO: Add to from_config constructor

    def transform_observation_space(
        self,
        observation_space: spaces.Dict,
    ):
        size = self._size
        observation_space = copy.deepcopy(observation_space)
        if size:
            for key in observation_space.spaces:
                if (
                    key in self.trans_keys
                    and observation_space.spaces[key].shape[-3:-1] != size
                ):
                    h, w = get_image_height_width(
                        observation_space.spaces[key], channels_last=True
                    )
                    logger.info(
                        "Center cropping observation size of %s from %s to %s"
                        % (key, (h, w), size)
                    )

                    observation_space.spaces[key] = overwrite_gym_box_shape(
                        observation_space.spaces[key], size
                    )
        return observation_space

    def _transform_obs(self, obs: torch.Tensor) -> torch.Tensor:
        return center_crop(
            obs,
            self._size,
            channels_last=self.channels_last,
        )

    @torch.no_grad()
    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        if self._size is not None:
            observations.update(
                {
                    sensor: self._transform_obs(observations[sensor])
                    for sensor in self.trans_keys
                    if sensor in observations
                }
            )
        return observations

    @classmethod
    def from_config(cls, config: Config):
        cc_config = config.RL.POLICY.OBS_TRANSFORMS.CENTER_CROPPER
        return cls(
            (
                cc_config.HEIGHT,
                cc_config.WIDTH,
            )
        )


class _DepthFrom(Enum):
    Z_VAL = 0
    OPTI_CENTER = 1


class CameraProjection(metaclass=abc.ABCMeta):
    """This is the base CameraProjection class that converts
    projection model of images into different one. This can be used for
    conversion between cubemap, equirect, fisheye images, etc.
    projection that project 3D points onto the image plane and
    unprojection that project image points onto unit sphere
    must be implemented."""

    def __init__(
        self,
        img_h: int,
        img_w: int,
        R: Optional[torch.Tensor] = None,
        depth_from: _DepthFrom = _DepthFrom.OPTI_CENTER,
    ):
        """Args:
        img_h: (int) the height of camera image
        img_w: (int) the width of camera image
        R: (torch.Tensor) 3x3 rotation matrix of camera
        depth_from: (_DepthFrom) the depth from z value or optical center
        """
        self.img_h = img_h
        self.img_w = img_w
        self.depth_from = depth_from

        # Camera rotation: points in world coord = R @ points in camera coord
        if R is not None:
            self.R = R.float()
        else:
            self.R = None

    @abc.abstractmethod
    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Project points in world coord onto image planes.
        Args:
            world_pts: 3D points in world coord
        Returns:
            proj_pts: Projected points for grid_sample, -1 <= proj_pts <= 1
            valid_mask: True if the point is valid (inside FoV)
        """

    @abc.abstractmethod
    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Unproject 2D image points onto unit sphere.
        Args:
            with_rotation: If True, unprojected points is in world coord.
                           If False, unprojected points is in camera coord.
        Returns:
            unproj_pts: Unprojected 3D points on unit sphere
            valid_mask: True if the point is valid (inside FoV)
        """

    @property
    def rotation(self):
        """Camera rotation: points in world coord = R @ points in camera coord"""
        if self.R is None:
            return torch.eye(3, dtype=torch.float32)
        else:
            return self.R

    @property
    def shape(self):
        """Camera image shape: (img_h, img_w)"""
        return (self.img_h, self.img_w)

    def size(self):
        """Camera image shape: (img_h, img_w)"""
        return self.shape

    def camcoord2worldcoord(self, pts: torch.Tensor):
        """Convert points in camera coords into points in world coords.
        Args:
            pts: 3D points in camera coords
        Returns:
            rotated_pts: 3D points in world coords
        """
        if self.R is None:
            return pts
        else:
            # Rotate points according to camera rotation
            _h, _w, _ = pts.shape
            # points in world coord = R @ points in camera coord
            rotated_pts = torch.matmul(pts.view((-1, 3)), self.R.T)
            return rotated_pts.view(_h, _w, 3)

    def worldcoord2camcoord(self, pts: torch.Tensor):
        """Convert points in world coords into points in camera coords.
        Args:
            pts: 3D points in world coords
        Returns:
            rotated_pts: 3D points in camera coords
        """
        if self.R is None:
            return pts
        else:
            # Rotate points according to camera rotation
            _h, _w, _ = pts.shape
            # points in camera coord = R.T @ points in world coord
            rotated_pts = torch.matmul(pts.view((-1, 3)), self.R)
            return rotated_pts.view(_h, _w, 3)


class PerspectiveProjection(CameraProjection):
    """This is the perspective camera projection class."""

    def __init__(
        self,
        img_h: int,
        img_w: int,
        f: Optional[float] = None,
        R: Optional[torch.Tensor] = None,
    ):
        """Args:
        img_h: (int) the height of camera image
        img_w: (int) the width of camera image
        f: (float) the focal length of camera
        R: (torch.Tensor) 3x3 rotation matrix of camera
        """
        super(PerspectiveProjection, self).__init__(
            img_h, img_w, R, _DepthFrom.Z_VAL
        )
        if f is None:
            self.f = max(img_h, img_w) / 2
        else:
            self.f = f

    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Rotate world points according to camera rotation
        world_pts = self.worldcoord2camcoord(world_pts)

        # Project points onto image plane
        img_pts = self.f * world_pts / torch.abs(world_pts[..., 2:3])
        cx = self.img_w / 2
        cy = self.img_h / 2
        u = img_pts[..., 0] + cx
        v = img_pts[..., 1] + cy

        # For grid_sample, -1 <= proj_pts <= 1
        mapx = 2 * u / self.img_w - 1.0
        mapy = 2 * v / self.img_h - 1.0
        proj_pts = torch.stack([mapx, mapy], dim=-1)

        # Valid mask
        valid_mask = torch.abs(proj_pts).max(-1)[0] <= 1  # -1 <= grid.xy <= 1
        valid_mask *= img_pts[..., 2] > 0
        return proj_pts, valid_mask

    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        v, u = torch.meshgrid(
            torch.arange(self.img_h), torch.arange(self.img_w)
        )
        x = (u + 0.5) - self.img_w / 2
        y = (v + 0.5) - self.img_h / 2
        z = torch.full_like(x, self.f, dtype=torch.float)
        unproj_pts = torch.stack([x, y, z], dim=-1)
        # Project on unit shpere
        unproj_pts /= torch.norm(unproj_pts, dim=-1, keepdim=True)
        # All points in image are valid
        valid_mask = torch.full(unproj_pts.shape[:2], True, dtype=torch.bool)

        # Rotate unproj_pts points according to camera rotation
        if with_rotation:
            unproj_pts = self.camcoord2worldcoord(unproj_pts)

        return unproj_pts, valid_mask


class EquirectProjection(CameraProjection):
    """This is the equirectanglar camera projection class."""

    def __init__(
        self, img_h: int, img_w: int, R: Optional[torch.Tensor] = None
    ):
        """Args:
        img_h: (int) the height of equirectanglar camera image
        img_w: (int) the width of equirectanglar camera image
        R: (torch.Tensor) 3x3 rotation matrix of camera
        """
        super(EquirectProjection, self).__init__(img_h, img_w, R)

    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Rotate world points according to camera rotation
        world_pts = self.worldcoord2camcoord(world_pts)

        x, y, z = world_pts[..., 0], world_pts[..., 1], world_pts[..., 2]
        # x,y,z to theta, phi
        theta = torch.atan2(x, z)
        c = torch.sqrt(x * x + z * z)
        phi = torch.atan2(y, c)

        # For grid_sample, -1 <= proj_pts <= 1
        mapx = theta / np.pi
        mapy = phi / (np.pi / 2)
        proj_pts = torch.stack([mapx, mapy], dim=-1)

        # All points in image are valid
        valid_mask = torch.full(proj_pts.shape[:2], True, dtype=torch.bool)
        return proj_pts, valid_mask

    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        theta_map, phi_map = self.get_theta_phi_map(self.img_h, self.img_w)
        unproj_pts = self.angle2sphere(theta_map, phi_map)
        # All points in image are valid
        valid_mask = torch.full(unproj_pts.shape[:2], True, dtype=torch.bool)
        # Rotate unproj_pts points according to camera rotation
        if with_rotation:
            unproj_pts = self.camcoord2worldcoord(unproj_pts)
        return unproj_pts, valid_mask

    def get_theta_phi_map(
        self, img_h: int, img_w: int
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Get theta and phi map for equirectangular image.
        PI < theta_map < PI,  PI/2 < phi_map < PI/2
        """
        phi, theta = torch.meshgrid(torch.arange(img_h), torch.arange(img_w))
        theta_map = (theta + 0.5) * 2 * np.pi / img_w - np.pi
        phi_map = (phi + 0.5) * np.pi / img_h - np.pi / 2
        return theta_map, phi_map

    def angle2sphere(
        self, theta_map: torch.Tensor, phi_map: torch.Tensor
    ) -> torch.Tensor:
        """Project points on unit sphere based on theta and phi map."""
        sin_theta = torch.sin(theta_map)
        cos_theta = torch.cos(theta_map)
        sin_phi = torch.sin(phi_map)
        cos_phi = torch.cos(phi_map)
        return torch.stack(
            [cos_phi * sin_theta, sin_phi, cos_phi * cos_theta], dim=-1
        )


class FisheyeProjection(CameraProjection):
    r"""This is the fisheye camera projection class.
    The camera model is based on the Double Sphere Camera Model (Usenko et. al.;3DV 2018).
    Paper: https://arxiv.org/abs/1807.08957
    Implementation: https://github.com/matsuren/dscamera
    """

    def __init__(
        self,
        img_h: int,
        img_w: int,
        fish_fov: float,
        cx: float,
        cy: float,
        fx: float,
        fy: float,
        xi: float,
        alpha: float,
        R: Optional[torch.Tensor] = None,
    ):
        """Args:
        img_h: (int) the height of fisheye camera image
        img_w: (int) the width of fisheye camera image
        fish_fov: (float) the fov of fisheye camera in degrees
        cx, cy: (float) the optical center of the fisheye camera
        fx, fy, xi, alpha: (float) the fisheye camera model parameters
        R: (torch.Tensor) 3x3 rotation matrix of camera
        """
        super(FisheyeProjection, self).__init__(img_h, img_w, R)

        self.fish_fov = fish_fov  # FoV in degrees
        fov_rad = self.fish_fov / 180 * np.pi  # FoV in radians
        self.fov_cos = np.cos(fov_rad / 2)
        self.fish_param = [cx, cy, fx, fy, xi, alpha]

    def projection(
        self, world_pts: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Rotate world points according to camera rotation
        world_pts = self.worldcoord2camcoord(world_pts)

        # Unpack parameters
        cx, cy, fx, fy, xi, alpha = self.fish_param
        # Unpack 3D world points
        x, y, z = world_pts[..., 0], world_pts[..., 1], world_pts[..., 2]

        # Calculate fov
        world_pts_fov_cos = z  # point3D @ z_axis
        fov_mask = world_pts_fov_cos >= self.fov_cos

        # Calculate projection
        x2 = x * x
        y2 = y * y
        z2 = z * z
        d1 = torch.sqrt(x2 + y2 + z2)
        zxi = xi * d1 + z
        d2 = torch.sqrt(x2 + y2 + zxi * zxi)

        div = alpha * d2 + (1 - alpha) * zxi
        u = fx * x / div + cx
        v = fy * y / div + cy

        # Projected points on image plane
        # For grid_sample, -1 <= proj_pts <= 1
        mapx = 2 * u / self.img_w - 1.0
        mapy = 2 * v / self.img_h - 1.0
        proj_pts = torch.stack([mapx, mapy], dim=-1)

        # Check valid area
        if alpha <= 0.5:
            w1 = alpha / (1 - alpha)
        else:
            w1 = (1 - alpha) / alpha
        w2 = w1 + xi / np.sqrt(2 * w1 * xi + xi * xi + 1)
        valid_mask = z > -w2 * d1
        valid_mask *= fov_mask

        return proj_pts, valid_mask

    def unprojection(
        self, with_rotation: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        # Unpack parameters
        cx, cy, fx, fy, xi, alpha = self.fish_param

        # Calculate unprojection
        v, u = torch.meshgrid(
            [torch.arange(self.img_h), torch.arange(self.img_w)]
        )
        mx = (u - cx) / fx
        my = (v - cy) / fy
        r2 = mx * mx + my * my
        mz = (1 - alpha * alpha * r2) / (
            alpha * torch.sqrt(1 - (2 * alpha - 1) * r2) + 1 - alpha
        )
        mz2 = mz * mz

        k1 = mz * xi + torch.sqrt(mz2 + (1 - xi * xi) * r2)
        k2 = mz2 + r2
        k = k1 / k2

        # Unprojected unit vectors
        unproj_pts = k.unsqueeze(-1) * torch.stack([mx, my, mz], dim=-1)
        unproj_pts[..., 2] -= xi

        # Calculate fov
        unproj_fov_cos = unproj_pts[..., 2]  # unproj_pts @ z_axis
        fov_mask = unproj_fov_cos >= self.fov_cos
        if alpha > 0.5:
            fov_mask *= r2 <= (1 / (2 * alpha - 1))

        # Rotate unproj_pts points according to camera rotation
        if with_rotation:
            unproj_pts = self.camcoord2worldcoord(unproj_pts)

        return unproj_pts, fov_mask


class ProjectionConverter(nn.Module):
    r"""This is the implementation to convert {cubemap, equirect, fisheye} images
    into {perspective, equirect, fisheye} images.
    """

    def __init__(
        self,
        input_projections: Union[List[CameraProjection], CameraProjection],
        output_projections: Union[List[CameraProjection], CameraProjection],
    ):
        """Args:
        input_projections: input images of projection models
        output_projections: generated image of projection models
        """
        super(ProjectionConverter, self).__init__()
        # Convert to list
        if not isinstance(input_projections, list):
            input_projections = [input_projections]
        if not isinstance(output_projections, list):
            output_projections = [output_projections]

        self.input_models = input_projections
        self.output_models = output_projections
        self.input_len = len(self.input_models)
        self.output_len = len(self.output_models)

        # Check image size
        input_size = self.input_models[0].size()
        for it in self.input_models:
            assert (
                input_size == it.size()
            ), "All input models must have the same image size"

        output_size = self.output_models[0].size()
        for it in self.output_models:
            assert (
                output_size == it.size()
            ), "All output models must have the same image size"

        # Check if depth conversion is required
        # If depth is in z value in input, conversion is required
        self.input_zfactor = self.calculate_zfactor(self.input_models)
        # If depth is in z value in output, inverse conversion is required
        self.output_zfactor = self.calculate_zfactor(
            self.output_models, inverse=True
        )

        # grids shape: (output_len, input_len, output_img_h, output_img_w, 2)
        self.grids = self.generate_grid()
        # _grids_cache shape: (batch_size*output_len*input_len, output_img_h, output_img_w, 2)
        self._grids_cache = None

    def _generate_grid_one_output(
        self, output_model: CameraProjection
    ) -> torch.Tensor:
        # Obtain points on unit sphere
        world_pts, not_assigned_mask = output_model.unprojection()
        # Generate grid
        grids = []
        for input_model in self.input_models:
            grid, input_mask = input_model.projection(world_pts)
            # Make sure each point is only assigned to single input
            input_mask *= not_assigned_mask
            # Values bigger than one will be ignored by grid_sample
            grid[~input_mask] = 2
            # Update not_assigned_mask
            not_assigned_mask *= ~input_mask
            grids.append(grid)
        grids = torch.stack(grids, dim=0)
        return grids

    def generate_grid(self) -> torch.Tensor:
        multi_output_grids = []
        for output_model in self.output_models:
            grids = self._generate_grid_one_output(output_model)
            multi_output_grids.append(grids.unsqueeze(1))
        multi_output_grids = torch.cat(multi_output_grids, dim=1)
        return multi_output_grids  # input_len, output_len, output_img_h, output_img_w, 2

    def _convert(self, batch: torch.Tensor) -> torch.Tensor:
        """Takes a batch of images stacked in proper order and converts thems,
        reduces batch size by input_len."""
        batch_size, ch, _H, _W = batch.shape
        out_h, out_w = self.output_models[0].size()
        if batch_size == 0 or batch_size % self.input_len != 0:
            raise ValueError(f"Batch size should be {self.input_len}x")
        output = torch.nn.functional.grid_sample(
            batch,
            self._grids_cache,
            align_corners=True,
            padding_mode="zeros",
        )
        output = output.view(
            batch_size // self.input_len,
            self.input_len,
            ch,
            out_h,
            out_w,
        ).sum(dim=1)
        return output  # output_len * batch_size, ch, output_model.img_h, output_model.img_w

    def to_converted_tensor(self, batch: torch.Tensor) -> torch.Tensor:
        """Convert tensors based on projection models. If there are two
        batches from two envs (R_1st, G_1st, B_1st) and (R_2nd, G_2nd, B_2nd),
        the input order is [R_1st, G_1st, B_1st, R_2nd, G_2nd, B_2nd]
        """
        # batch tensor order should be NCHW
        batch_size, ch, in_h, in_w = batch.size()

        out_h, out_w = self.output_models[0].size()

        # Check whether batch size is len(self.input_models) x
        if batch_size == 0 or batch_size % self.input_len != 0:
            raise ValueError(f"Batch size should be {self.input_len}x")

        # How many sets of input.
        num_input_set = batch_size // self.input_len

        # to(device) is a NOOP after the first call
        self.grids = self.grids.to(batch.device)

        # Adjust batch for multiple outputs
        # batch must be [1st batch * output_len, 2nd batch * output_len, ...]
        # not that [1st batch, 2nd batch, ...] * output_len
        multi_out_batch = (
            batch.view(num_input_set, self.input_len, ch, in_h, in_w)
            .repeat(1, self.output_len, 1, 1, 1)
            .view(self.output_len * batch_size, ch, in_h, in_w)
        )

        # Cache the repeated grids for subsequent batches
        if (
            self._grids_cache is None
            or self._grids_cache.size()[0] != multi_out_batch.size()[0]
        ):
            # batch size is more than one
            self._grids_cache = self.grids.repeat(
                num_input_set, 1, 1, 1, 1
            ).view(batch_size * self.output_len, out_h, out_w, 2)
        self._grids_cache = self._grids_cache.to(batch.device)

        return self._convert(multi_out_batch)

    def calculate_zfactor(
        self, projections: List[CameraProjection], inverse: bool = False
    ) -> Optional[torch.Tensor]:
        """Calculate z factor based on camera projection models. z_factor is
        used for converting depth in z value to depth from optical center
        (for input_models) or conversion of depth from optical center to depth
        in z value (inverse = True, for output_models). Whether the conversion
        is required or not is decided based on depth_from property of
        CameraProjection class.
        Args:
            projections: input or output projection models
            inverse: True to convert depth from optical center to z value
                     False to convert z value to depth from optical center
        Returns:
            z_factors: z factor. Return None if conversion is not required.
        """
        z_factors = []
        for cam in projections:
            if cam.depth_from == _DepthFrom.Z_VAL:
                pts_on_sphere, _ = cam.unprojection(with_rotation=False)
                zval_to_optcenter = 1 / pts_on_sphere[..., 2]
                z_factors.append(zval_to_optcenter.unsqueeze(0))
            else:
                all_one = torch.full(
                    (1, cam.img_h, cam.img_w), 1.0, dtype=torch.float
                )
                z_factors.append(all_one)
        z_factors = torch.stack(z_factors)

        if (z_factors == 1.0).all():
            # All input cameras have depth from optical center
            return None
        else:
            if not inverse:
                # for input_models
                return z_factors
            else:
                # for output_models
                return 1 / z_factors

    def forward(
        self, batch: torch.Tensor, is_depth: bool = False
    ) -> torch.Tensor:

        # Depth conversion for input tensors
        if is_depth and self.input_zfactor is not None:
            input_b = batch.size()[0] // self.input_len
            self.input_zfactor = self.input_zfactor.to(batch.device)
            batch = batch * self.input_zfactor.repeat(input_b, 1, 1, 1)

        # Common operator to convert projection models
        out = self.to_converted_tensor(batch)

        # Depth conversion for output tensors
        if is_depth and self.output_zfactor is not None:
            output_b = out.size()[0] // self.output_len
            self.output_zfactor = self.output_zfactor.to(batch.device)
            out = out * self.output_zfactor.repeat(output_b, 1, 1, 1)

        return out


def get_cubemap_projections(
    img_h: int = 256, img_w: int = 256
) -> List[CameraProjection]:
    """Get cubemap camera projections that consist of six PerspectiveCameras.
    The orders are 'BACK', 'DOWN', 'FRONT', 'LEFT', 'RIGHT', 'UP'.
    Args:
        img_h: (int) the height of camera image
        img_w: (int) the width of camera image

    The rotation matrices are equivalent to
    .. code-block:: python
        from scipy.spatial.transform import Rotation
        rotations = [
            Rotation.from_euler("y", 180, degrees=True),  # Back
            Rotation.from_euler("x", -90, degrees=True),  # Down
            Rotation.from_euler("x", 0, degrees=True),  # Front
            Rotation.from_euler("y", -90, degrees=True),  # Left
            Rotation.from_euler("y", 90, degrees=True),  # Right
            Rotation.from_euler("x", 90, degrees=True)  # Up
        ]
    """
    rotations = [
        torch.tensor([[-1, 0, 0], [0, 1, 0], [0, 0, -1]]),  # Back
        torch.tensor([[1, 0, 0], [0, 0, 1], [0, -1, 0]]),  # Down
        torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]]),  # Front
        torch.tensor([[0, 0, -1], [0, 1, 0], [1, 0, 0]]),  # Left
        torch.tensor([[0, 0, 1], [0, 1, 0], [-1, 0, 0]]),  # Right
        torch.tensor([[1, 0, 0], [0, 0, -1], [0, 1, 0]]),  # Up
    ]

    projections = []
    for rot in rotations:
        cam = PerspectiveProjection(img_h, img_w, R=rot)
        projections.append(cam)
    return projections


class Cube2Equirect(ProjectionConverter):
    """This is the backend Cube2Equirect nn.module that does the stiching.
    Inspired from https://github.com/fuenwang/PanoramaUtility and
    optimized for modern PyTorch."""

    def __init__(self, equ_h: int, equ_w: int):
        """Args:
        equ_h: (int) the height of the generated equirect
        equ_w: (int) the width of the generated equirect
        """

        # Cubemap input
        input_projections = get_cubemap_projections()

        # Equirectangular output
        output_projection = EquirectProjection(equ_h, equ_w)
        super(Cube2Equirect, self).__init__(
            input_projections, output_projection
        )


class ProjectionTransformer(ObservationTransformer):
    r"""
    ProjectionTransformer base class. It can be used to  convert {cubemap, equirect, fisheye} images
    into {perspective, equirect, fisheye} images in ObservationTransformer.
    """

    def __init__(
        self,
        converter: ProjectionConverter,
        sensor_uuids: List[str],
        image_shape: Tuple[int, int],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param converter: ProjectionConverter class
        :param sensor_uuids: List of sensor_uuids
        :param image_shape: The shape of the output image (height, width)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """
        super(ProjectionTransformer, self).__init__()
        num_sensors = len(sensor_uuids)
        assert (
            num_sensors % converter.input_len == 0 and num_sensors != 0
        ), f"{len(sensor_uuids)}: length of sensors is not a multiple of {converter.input_len}"
        # TODO verify attributes of the sensors in the config if possible. Think about API design
        assert (
            len(image_shape) == 2
        ), f"image_shape must be a tuple of (height, width), given: {image_shape}"
        self.sensor_uuids: List[str] = sensor_uuids
        self.img_shape: Tuple[int, int] = image_shape
        self.channels_last: bool = channels_last
        self.converter = converter
        if target_uuids is None:
            self.target_uuids: List[str] = self.sensor_uuids[::6]
        else:
            self.target_uuids: List[str] = target_uuids
        self.depth_key = depth_key

    def transform_observation_space(
        self,
        observation_space: spaces.Dict,
    ):
        r"""Transforms the target UUID's sensor obs_space so it matches the new shape (H, W)"""
        # Transforms the observation space to of the target UUID
        for i, key in enumerate(self.target_uuids):
            assert (
                key in observation_space.spaces
            ), f"{key} not found in observation space: {observation_space.spaces}"
            h, w = get_image_height_width(
                observation_space.spaces[key], channels_last=True
            )
            in_len = self.converter.input_len
            logger.info(
                f"Overwrite sensor: {key} from size of ({h}, {w}) to image of"
                f" {self.img_shape} from sensors: {self.sensor_uuids[i*in_len:(i+1)*in_len]}"
            )
            if (h, w) != self.img_shape:
                observation_space.spaces[key] = overwrite_gym_box_shape(
                    observation_space.spaces[key], self.img_shape
                )
        return observation_space

    @torch.no_grad()
    def forward(
        self, observations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:

        for i, target_sensor_uuid in enumerate(self.target_uuids):
            # number of input and input sensor uuids
            in_len = self.converter.input_len
            in_sensor_uuids = self.sensor_uuids[i * in_len : (i + 1) * in_len]

            # If the sensor is depth
            is_depth = any(self.depth_key in s for s in in_sensor_uuids)

            # The UUID we are overwriting
            assert target_sensor_uuid in in_sensor_uuids
            sensor_obs = [observations[sensor] for sensor in in_sensor_uuids]
            target_obs = observations[target_sensor_uuid]
            sensor_dtype = target_obs.dtype
            # Stacking along axis makes the flattening go in the right order.
            imgs = torch.stack(sensor_obs, axis=1)
            imgs = torch.flatten(imgs, end_dim=1)
            if not self.channels_last:
                imgs = imgs.permute((0, 3, 1, 2))  # NHWC => NCHW
            imgs = imgs.float()  # NCHW
            # Here is where the projection conversion happens
            output = self.converter(imgs, is_depth=is_depth)

            # for debugging
            # torchvision.utils.save_image(output, f'sample_eqr_{target_sensor_uuid}.jpg', normalize=True, range=(0, 255) if 'rgb' in target_sensor_uuid else (0, 1))
            output = output.to(dtype=sensor_dtype)
            if not self.channels_last:
                output = output.permute((0, 2, 3, 1))  # NCHW => NHWC
            observations[target_sensor_uuid] = output
        return observations


@baseline_registry.register_obs_transformer()
class CubeMap2Equirect(ProjectionTransformer):
    r"""This is an experimental use of ObservationTransformer that converts a cubemap
    output to an equirectangular one through projection. This needs to be fed
    a list of 6 cameras at various orientations but will be able to stitch a
    360 sensor out of these inputs. The code below will generate a config that
    has the 6 sensors in the proper orientations. This code also assumes a 90
    FOV.

    Sensor order for cubemap stiching is Back, Down, Front, Left, Right, Up.
    The output will be writen the UUID of the first sensor.
    """

    def __init__(
        self,
        sensor_uuids: List[str],
        eq_shape: Tuple[int, int],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param sensor_uuids: List of sensor_uuids: Back, Down, Front, Left, Right, Up.
        :param eq_shape: The shape of the equirectangular output (height, width)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """

        converter = Cube2Equirect(eq_shape[0], eq_shape[1])
        super(CubeMap2Equirect, self).__init__(
            converter,
            sensor_uuids,
            eq_shape,
            channels_last,
            target_uuids,
            depth_key,
        )

    @classmethod
    def from_config(cls, config):
        cube2eq_config = config.RL.POLICY.OBS_TRANSFORMS.CUBE2EQ
        if hasattr(cube2eq_config, "TARGET_UUIDS"):
            # Optional Config Value to specify target UUID
            target_uuids = cube2eq_config.TARGET_UUIDS
        else:
            target_uuids = None
        return cls(
            cube2eq_config.SENSOR_UUIDS,
            eq_shape=(
                cube2eq_config.HEIGHT,
                cube2eq_config.WIDTH,
            ),
            target_uuids=target_uuids,
        )


class Cube2Fisheye(ProjectionConverter):
    r"""This is the implementation to generate fisheye images from cubemap images.
    The camera model is based on the Double Sphere Camera Model (Usenko et. al.;3DV 2018).
    Paper: https://arxiv.org/abs/1807.08957
    """

    def __init__(
        self,
        fish_h: int,
        fish_w: int,
        fish_fov: float,
        cx: float,
        cy: float,
        fx: float,
        fy: float,
        xi: float,
        alpha: float,
    ):
        """Args:
        fish_h: (int) the height of the generated fisheye
        fish_w: (int) the width of the generated fisheye
        fish_fov: (float) the fov of the generated fisheye in degrees
        cx, cy: (float) the optical center of the generated fisheye
        fx, fy, xi, alpha: (float) the fisheye camera model parameters
        """

        # Cubemap input
        input_projections = get_cubemap_projections()

        # Fisheye output
        output_projection = FisheyeProjection(
            fish_h, fish_w, fish_fov, cx, cy, fx, fy, xi, alpha
        )
        super(Cube2Fisheye, self).__init__(
            input_projections, output_projection
        )


@baseline_registry.register_obs_transformer()
class CubeMap2Fisheye(ProjectionTransformer):
    r"""This is an experimental use of ObservationTransformer that converts a cubemap
    output to a fisheye one through projection. This needs to be fed
    a list of 6 cameras at various orientations but will be able to stitch a
    fisheye image out of these inputs. The code below will generate a config that
    has the 6 sensors in the proper orientations. This code also assumes a 90
    FOV.

    Sensor order for cubemap stiching is Back, Down, Front, Left, Right, Up.
    The output will be writen the UUID of the first sensor.
    """

    def __init__(
        self,
        sensor_uuids: List[str],
        fish_shape: Tuple[int, int],
        fish_fov: float,
        fish_params: Tuple[float],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param sensor_uuids: List of sensor_uuids: Back, Down, Front, Left, Right, Up.
        :param fish_shape: The shape of the fisheye output (height, width)
        :param fish_fov: The FoV of the fisheye output in degrees
        :param fish_params: The camera parameters of fisheye output (f, xi, alpha)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """

        assert (
            len(fish_params) == 3
        ), "fish_params must have three parameters (f, xi, alpha)"
        # fisheye camera parameters
        fx = fish_params[0] * min(fish_shape)
        fy = fx
        cx = fish_shape[1] / 2
        cy = fish_shape[0] / 2
        xi = fish_params[1]
        alpha = fish_params[2]
        converter: ProjectionConverter = Cube2Fisheye(
            fish_shape[0], fish_shape[1], fish_fov, cx, cy, fx, fy, xi, alpha
        )

        super(CubeMap2Fisheye, self).__init__(
            converter,
            sensor_uuids,
            fish_shape,
            channels_last,
            target_uuids,
            depth_key,
        )

    @classmethod
    def from_config(cls, config):
        cube2fish_config = config.RL.POLICY.OBS_TRANSFORMS.CUBE2FISH
        if hasattr(cube2fish_config, "TARGET_UUIDS"):
            # Optional Config Value to specify target UUID
            target_uuids = cube2fish_config.TARGET_UUIDS
        else:
            target_uuids = None
        return cls(
            cube2fish_config.SENSOR_UUIDS,
            fish_shape=(
                cube2fish_config.HEIGHT,
                cube2fish_config.WIDTH,
            ),
            fish_fov=cube2fish_config.FOV,
            fish_params=cube2fish_config.PARAMS,
            target_uuids=target_uuids,
        )


class Equirect2Cube(ProjectionConverter):
    """This is the backend Equirect2CubeMap that converts equirectangular image
    to cubemap images."""

    def __init__(self, img_h: int, img_w: int):
        """Args:
        img_h: (int) the height of the generated cubemap
        img_w: (int) the width of the generated cubemap
        """

        # Equirectangular input
        input_projection = EquirectProjection(256, 512)

        #  Cubemap output
        output_projections = get_cubemap_projections(img_h, img_w)
        super(Equirect2Cube, self).__init__(
            input_projection, output_projections
        )


@baseline_registry.register_obs_transformer()
class Equirect2CubeMap(ProjectionTransformer):
    r"""This is an experimental use of ObservationTransformer that converts
    an equirectangular image to cubemap images.
    Cubemap order is Back, Down, Front, Left, Right, Up.
    The output will be writen the UUID of the first sensor.
    """

    def __init__(
        self,
        sensor_uuids: List[str],
        img_shape: Tuple[int, int],
        channels_last: bool = False,
        target_uuids: Optional[List[str]] = None,
        depth_key: str = "depth",
    ):
        r""":param sensor_uuids: List of sensor_uuids: Back, Down, Front, Left, Right, Up.
        :param img_shape: The shape of the equirectangular output (height, width)
        :param channels_last: Are the channels last in the input
        :param target_uuids: Optional List of which of the sensor_uuids to overwrite
        :param depth_key: If sensor_uuids has depth_key substring, they are processed as depth
        """

        converter = Equirect2Cube(img_shape[0], img_shape[1])
        super(Equirect2CubeMap, self).__init__(
            converter,
            sensor_uuids,
            img_shape,
            channels_last,
            target_uuids,
            depth_key,
        )

    @classmethod
    def from_config(cls, config):
        eq2cube_config = config.RL.POLICY.OBS_TRANSFORMS.EQ2CUBE

        if hasattr(eq2cube_config, "TARGET_UUIDS"):
            # Optional Config Value to specify target UUID
            target_uuids = eq2cube_config.TARGET_UUIDS
        else:
            target_uuids = None
        return cls(
            eq2cube_config.SENSOR_UUIDS,
            img_shape=(
                eq2cube_config.HEIGHT,
                eq2cube_config.WIDTH,
            ),
            target_uuids=target_uuids,
        )


def get_active_obs_transforms(config: Config) -> List[ObservationTransformer]:
    active_obs_transforms = []
    if hasattr(config.RL.POLICY, "OBS_TRANSFORMS"):
        obs_transform_names = (
            config.RL.POLICY.OBS_TRANSFORMS.ENABLED_TRANSFORMS
        )
        for obs_transform_name in obs_transform_names:
            obs_trans_cls = baseline_registry.get_obs_transformer(
                obs_transform_name
            )
            obs_transform = obs_trans_cls.from_config(config)
            active_obs_transforms.append(obs_transform)
    return active_obs_transforms


def apply_obs_transforms_batch(
    batch: Dict[str, torch.Tensor],
    obs_transforms: Iterable[ObservationTransformer],
) -> Dict[str, torch.Tensor]:
    for obs_transform in obs_transforms:
        batch = obs_transform(batch)
    return batch


def apply_obs_transforms_obs_space(
    obs_space: spaces.Dict, obs_transforms: Iterable[ObservationTransformer]
) -> spaces.Dict:
    for obs_transform in obs_transforms:
        obs_space = obs_transform.transform_observation_space(obs_space)
    return obs_space
