# Copyright 2022 The Plenoptix Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Ray generator.
"""

from __future__ import annotations

import torch
from torch import nn
from torchtyping import TensorType

from nerfactory.cameras.rays import RayBundle
from nerfactory.datamanagers.structs import SceneBounds


class SceneBoundsCollider(nn.Module):
    """Module for setting near and far values for rays."""

    def __init__(self, **kwargs) -> None:
        self.kwargs = kwargs
        super().__init__()

    def forward(self, ray_bundle: RayBundle) -> RayBundle:
        """To be implemented."""
        raise NotImplementedError


class AABBBoxCollider(SceneBoundsCollider):
    """Module for colliding rays with the scene bounds to compute near and far values.

    Args:
        scene_bounds: scene bounds to apply to dataset
    """

    def __init__(self, scene_bounds: SceneBounds, near_plane: float = 0.0, **kwargs) -> None:
        super().__init__(**kwargs)
        self.scene_bounds = scene_bounds
        self.near_plane = near_plane

    def _intersect_with_aabb(
        self, rays_o: TensorType["num_rays", 3], rays_d: TensorType["num_rays", 3], aabb: TensorType[2, 3]
    ):
        """Returns collection of valid rays within a specified near/far bounding box along with a mask
        specifying which rays are valid

        Args:
            rays_o: (num_rays, 3) ray origins
            rays_d: (num_rays, 3) ray directions
            aabb: (2, 3) This is [min point (x,y,z), max point (x,y,z)]
        """
        # avoid divide by zero
        dir_fraction = 1.0 / (rays_d + 1e-6)

        # x
        t1 = (aabb[0, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]
        t2 = (aabb[1, 0] - rays_o[:, 0:1]) * dir_fraction[:, 0:1]
        # y
        t3 = (aabb[0, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]
        t4 = (aabb[1, 1] - rays_o[:, 1:2]) * dir_fraction[:, 1:2]
        # z
        t5 = (aabb[0, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]
        t6 = (aabb[1, 2] - rays_o[:, 2:3]) * dir_fraction[:, 2:3]

        nears = torch.max(
            torch.cat([torch.minimum(t1, t2), torch.minimum(t3, t4), torch.minimum(t5, t6)], dim=1), dim=1
        ).values
        fars = torch.min(
            torch.cat([torch.maximum(t1, t2), torch.maximum(t3, t4), torch.maximum(t5, t6)], dim=1), dim=1
        ).values

        # clamp to near plane
        nears = torch.clamp(nears, min=self.near_plane)

        # fars < 0: means the ray is behind the camera
        # nears > fars: means no intersection
        valid_mask = ((fars > nears).float() * (fars > 0).float()).bool()
        nears[~valid_mask] = 0.0
        fars[~valid_mask] = 0.0
        return nears, fars, valid_mask

    def forward(self, ray_bundle: RayBundle) -> RayBundle:
        """Intersects the rays with the scene bounds and updates the near and far values.
        Populates nears and fars fields and returns the ray_bundle.

        Args:
            ray_bundle: specified ray bundle to operate on
        """
        aabb = self.scene_bounds.aabb
        nears, fars, valid_mask = self._intersect_with_aabb(ray_bundle.origins, ray_bundle.directions, aabb)
        ray_bundle.nears = nears[..., None]
        ray_bundle.fars = fars[..., None]
        ray_bundle.valid_mask = valid_mask[..., None]
        return ray_bundle


class NearFarCollider(SceneBoundsCollider):
    """Sets the nears and fars with fixed values.

    Args:
        near_plane: distance to near plane
        far_plane: distance to far plane
    """

    def __init__(self, near_plane: float, far_plane: float, **kwargs) -> None:
        self.near_plane = near_plane
        self.far_plane = far_plane
        super().__init__(**kwargs)

    def forward(self, ray_bundle: RayBundle) -> RayBundle:
        ones = torch.ones_like(ray_bundle.origins[..., 0:1])
        ray_bundle.nears = ones * self.near_plane
        ray_bundle.fars = ones * self.far_plane
        ray_bundle.valid_mask = ones.bool()
        return ray_bundle