# 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.

"""
Collection of sampling strategies
"""

import math
from abc import abstractmethod
from typing import Callable, Optional, Tuple

import nerfacc  # pylint: disable=import-error
import torch
from torch import nn
from torchtyping import TensorType

from nerfactory.cameras.rays import Frustums, RayBundle, RaySamples
from nerfactory.fields.density_fields.density_grid import DensityGrid


class Sampler(nn.Module):
    """Generate Samples

    Args:
        num_samples: number of samples to take
        density_field: density grid specifying weighting of
    """

    def __init__(
        self, num_samples: int, density_field: Optional[DensityGrid] = None, weight_threshold: float = 0.01
    ) -> None:
        super().__init__()
        self.num_samples = num_samples
        self.density_field = density_field
        self.weight_threshold = weight_threshold

    @abstractmethod
    def generate_ray_samples(self) -> RaySamples:
        """Generate Ray Samples"""

    def forward(self, *args, **kwargs) -> RaySamples:
        """Generate ray samples with optional density filtering"""
        ray_samples = self.generate_ray_samples(*args, **kwargs)
        if self.density_field is not None:
            densities = self.density_field.get_densities(ray_samples)
            deltas = torch.clamp(ray_samples.frustums.ends - ray_samples.frustums.starts, min=1e-10)
            density_threshold = torch.clamp(self.weight_threshold / deltas, max=self.density_field.mean_density)
            valid_mask = densities > density_threshold
            ray_samples.set_valid_mask(valid_mask & ray_samples.valid_mask)
        return ray_samples


class SpacedSampler(Sampler):
    """Sample points according to a function.

    Args:
        num_samples: Number of samples per ray
        spacing_fn: Function that dictates sample spacing (ie `lambda x : x` is uniform).
        spacing_fn_inv: The inverse of spacing_fn.
        train_stratified: Use stratified sampling during training. Defults to True
        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
    """

    def __init__(
        self,
        num_samples: int,
        spacing_fn: Callable,
        spacing_fn_inv: Callable,
        train_stratified=True,
        density_field: Optional[DensityGrid] = None,
        weight_threshold: float = 1e-2,
    ) -> None:
        super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold)
        self.train_stratified = train_stratified
        self.spacing_fn = spacing_fn
        self.spacing_fn_inv = spacing_fn_inv

    @torch.no_grad()
    def generate_ray_samples(
        self,
        ray_bundle: Optional[RayBundle] = None,
        num_samples: Optional[int] = None,
    ) -> RaySamples:
        """Generates position samples accoring to spacing function.

        Args:
            ray_bundle: Rays to generate samples for
            num_samples: Number of samples per ray

        Returns:
            Positions and deltas for samples along a ray
        """
        assert ray_bundle is not None
        assert ray_bundle.nears is not None
        assert ray_bundle.fars is not None

        num_samples = num_samples or self.num_samples
        num_rays = ray_bundle.origins.shape[0]

        bins = torch.linspace(0.0, 1.0, num_samples + 1).to(ray_bundle.origins.device)[None, ...]  # [1, num_samples+1]

        s_near, s_far = (self.spacing_fn(x) for x in (ray_bundle.nears, ray_bundle.fars))
        bins = self.spacing_fn_inv(bins * s_far + (1 - bins) * s_near)  # [num_rays, num_samples+1]

        if self.train_stratified and self.training:
            t_rand = torch.rand((num_rays, num_samples + 1), dtype=bins.dtype, device=bins.device)
            bin_centers = (bins[..., 1:] + bins[..., :-1]) / 2.0
            bin_upper = torch.cat([bin_centers, bins[..., -1:]], -1)
            bin_lower = torch.cat([bins[..., :1], bin_centers], -1)
            bins = bin_lower + (bin_upper - bin_lower) * t_rand

        ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1, None], bin_ends=bins[..., 1:, None])

        return ray_samples


class UniformSampler(SpacedSampler):
    """Sample uniformly along a ray

    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
    """

    def __init__(
        self,
        num_samples: int,
        train_stratified=True,
        density_field: Optional[DensityGrid] = None,
        weight_threshold: float = 1e-2,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: x,
            spacing_fn_inv=lambda x: x,
            train_stratified=train_stratified,
            density_field=density_field,
            weight_threshold=weight_threshold,
        )


class LinearDisparitySampler(SpacedSampler):
    """Sample linearly in disparity along a ray

    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
    """

    def __init__(
        self,
        num_samples: int,
        train_stratified=True,
        density_field: Optional[DensityGrid] = None,
        weight_threshold: float = 1e-2,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=lambda x: 1 / x,
            spacing_fn_inv=lambda x: 1 / x,
            train_stratified=train_stratified,
            density_field=density_field,
            weight_threshold=weight_threshold,
        )


class SqrtSampler(SpacedSampler):
    """Square root sampler along a ray

    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
    """

    def __init__(
        self,
        num_samples: int,
        train_stratified=True,
        density_field: Optional[DensityGrid] = None,
        weight_threshold: float = 1e-2,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=torch.sqrt,
            spacing_fn_inv=lambda x: x**2,
            train_stratified=train_stratified,
            density_field=density_field,
            weight_threshold=weight_threshold,
        )


class LogSampler(SpacedSampler):
    """Log sampler along a ray

    Args:
        num_samples: Number of samples per ray
        train_stratified: Use stratified sampling during training. Defults to True
        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
    """

    def __init__(
        self,
        num_samples: int,
        train_stratified=True,
        density_field: Optional[DensityGrid] = None,
        weight_threshold: float = 1e-2,
    ) -> None:
        super().__init__(
            num_samples=num_samples,
            spacing_fn=torch.log,
            spacing_fn_inv=torch.exp,
            train_stratified=train_stratified,
            density_field=density_field,
            weight_threshold=weight_threshold,
        )


class PDFSampler(Sampler):
    """Sample based on probability distribution

    Args:
        num_samples: Number of samples per ray
        train_stratified: Randomize location within each bin during training. Defaults to True
        include_original: Add original samples to ray. Defaults to True
        density_field: Density grid. If provides, samples below weight_threshold as set as invalid.
        weight_threshold: Removes samples below threshold weight. Only used if density field is provided.
        histogram_padding: Amount to weights prior to computing PDF. Defaults to 0.01.
    """

    def __init__(
        self,
        num_samples: int,
        train_stratified: bool = True,
        include_original: bool = True,
        density_field: Optional[DensityGrid] = None,
        weight_threshold: float = 1e-2,
        histogram_padding: float = 0.01,
    ) -> None:
        super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold)
        self.train_stratified = train_stratified
        self.include_original = include_original
        self.histogram_padding = histogram_padding

    @torch.no_grad()
    def generate_ray_samples(
        self,
        ray_bundle: Optional[RayBundle] = None,
        ray_samples: Optional[RaySamples] = None,
        weights: TensorType[..., "num_samples", 1] = None,
        num_samples: Optional[int] = None,
        eps: float = 1e-5,
    ) -> RaySamples:
        """Generates position samples given a distribution.

        Args:
            ray_bundle: Rays to generate samples for
            ray_samples: Existing ray samples
            weights: Weights for each bin
            num_samples: Number of samples per ray
            eps: Small value to prevent numerical issues. Defaults to 1e-5

        Returns:
            Positions and deltas for samples along a ray
        """

        if ray_samples is None or ray_bundle is None:
            raise ValueError("ray_samples and ray_bundle must be provided")

        num_samples = num_samples or self.num_samples
        num_bins = num_samples + 1

        weights = weights[..., 0] + self.histogram_padding

        # Add small offset to rays with zero weight to prevent NaNs
        weights_sum = torch.sum(weights, dim=-1, keepdim=True)
        padding = torch.relu(eps - weights_sum)
        weights = weights + padding / weights.shape[-1]
        weights_sum += padding

        pdf = weights / weights_sum
        cdf = torch.min(torch.ones_like(pdf), torch.cumsum(pdf, dim=-1))
        cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1)

        if self.train_stratified and self.training:
            # Stratified samples between 0 and 1
            u = torch.linspace(0.0, 1.0 - (1.0 / num_bins), steps=num_bins, device=cdf.device)
            u = u.expand(size=(*cdf.shape[:-1], num_bins))
            u = u + torch.rand(size=(*cdf.shape[:-1], num_bins), device=cdf.device) / num_bins
        else:
            # Uniform samples between 0 and 1
            u = torch.linspace(0.0, 1.0, steps=num_bins, device=cdf.device)
            u = u.expand(size=(*cdf.shape[:-1], num_bins))
        u = u.contiguous()

        # Force bins to not have a gap between them. Kinda hacky, should reconsider.
        existing_bins = torch.cat(
            [
                ray_samples.frustums.starts[..., :1, 0],
                (ray_samples.frustums.starts[..., 1:, 0] + ray_samples.frustums.ends[..., :-1, 0]) / 2.0,
                ray_samples.frustums.ends[..., -1:, 0],
            ],
            dim=-1,
        )

        inds = torch.searchsorted(cdf, u, side="right")
        below = torch.clamp(inds - 1, 0, existing_bins.shape[-1] - 1)
        above = torch.clamp(inds, 0, existing_bins.shape[-1] - 1)
        cdf_g0 = torch.gather(cdf, -1, below)
        bins_g0 = torch.gather(existing_bins, -1, below)
        cdf_g1 = torch.gather(cdf, -1, above)
        bins_g1 = torch.gather(existing_bins, -1, above)

        t = torch.clip(torch.nan_to_num((u - cdf_g0) / (cdf_g1 - cdf_g0), 0), 0, 1)
        bins = bins_g0 + t * (bins_g1 - bins_g0)

        if self.include_original:
            bins, _ = torch.sort(torch.cat([existing_bins, bins], -1), -1)

        # Stop gradients
        bins = bins.detach()

        ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1, None], bin_ends=bins[..., 1:, None])

        return ray_samples


class VolumetricSampler(Sampler):
    """Sampler inspired by the one proposed in the Instant-NGP paper.

    This sampler does ray-box AABB test, ray samples generation and density check, all together.
    """

    def __init__(
        self,
        aabb: TensorType[2, 3],
        density_fn: Callable[[TensorType["N", 3]], TensorType["N"]],
        grid_resolution: int = 128,
        num_samples: int = 1024,
    ) -> None:
        """Init.

        Args:
            aabb: Bounding box of the scene.
            density_fn: Function that takes a tensor of points and returns a tensor of densities.
            grid_resolution: Resolution of the density grid.
            num_samples: Number of samples per ray.
        """
        super().__init__(num_samples=num_samples)
        self.aabb = aabb
        self.density_fn = density_fn

        self.render_step_size = ((self.aabb[3:] - self.aabb[:3]).max() * math.sqrt(3) / num_samples).item()

        # setup occupancy field with eval function
        def occ_eval_fn(x: torch.Tensor) -> torch.Tensor:
            """Evaluate occupancy given positions.

            Args:
                x: positions with shape (N, 3).
            Returns:
                occupancy values with shape (N, 1).
            """
            density_after_activation = density_fn(x)
            # those two are similar when density is small.
            # occupancy = 1.0 - torch.exp(-density_after_activation * render_step_size)
            occupancy = density_after_activation * self.render_step_size
            return occupancy

        self.occ_field = nerfacc.OccupancyField(occ_eval_fn=occ_eval_fn, aabb=aabb, resolution=grid_resolution)

    def generate_ray_samples(self) -> RaySamples:
        raise RuntimeError(
            "The VolumetricSampler fuses sample generation and density check together. Please call forward() directly."
        )

    # pylint: disable=arguments-differ
    def forward(
        self,
        ray_bundle: RayBundle,
        num_samples: Optional[int] = None,
        near_plane: float = 0.0,
        remove_occluded_samples: bool = True,
    ) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples", 2]]:
        """Generate ray samples in a bounding box.

        Args:
            ray_bundle: Rays to generate samples for
            num_samples: Number of samples per ray
            near_plane: Near plane for raymarching
            remove_occluded_samples: Whether to remove occluded samples. (requires evaluating the density function)

        Returns:
            a tuple of (ray_samples, packed_info, ray_indices)
            The ray_samples are packed, only storing the valid samples.
            The packed_info contains all the information to recover packed samples into unpacked mode for rendering.
            The ray_indices contains the indices of the rays that each sample belongs to.
        """

        if self.density_field is not None:
            raise ValueError("NGPSpacedSampler does not support a density field")

        if num_samples is not None:
            render_step_size = ((self.aabb[3:] - self.aabb[:3]).max() * math.sqrt(3) / num_samples).item()
        else:
            num_samples = self.num_samples
            render_step_size = self.render_step_size

        rays_o = ray_bundle.origins.contiguous()
        rays_d = ray_bundle.directions.contiguous()
        if ray_bundle.camera_indices is not None:
            camera_indices = ray_bundle.camera_indices.contiguous()
        else:
            camera_indices = None

        packed_info, starts, ends = nerfacc.volumetric_marching(
            rays_o=rays_o,
            rays_d=rays_d,
            aabb=self.aabb,
            scene_resolution=self.occ_field.resolution,
            scene_occ_binary=self.occ_field.occ_grid_binary,
            render_step_size=render_step_size,
            near_plane=near_plane,
            stratified=self.training,
        )

        ray_indices = nerfacc.unpack_to_ray_indices(packed_info).long()
        origins = rays_o[ray_indices]
        dirs = rays_d[ray_indices]
        if camera_indices is not None:
            camera_indices = camera_indices[ray_indices]

        if remove_occluded_samples:
            with torch.no_grad():
                positions = origins + dirs * (starts + ends) / 2.0
                densities = self.density_fn(positions)
                if camera_indices is not None:
                    packed_info, starts, ends, origins, dirs, camera_indices = nerfacc.volumetric_rendering_steps(
                        packed_info, densities, starts, ends, origins, dirs, camera_indices
                    )
                else:
                    packed_info, starts, ends, origins, dirs = nerfacc.volumetric_rendering_steps(
                        packed_info, densities, starts, ends, origins, dirs
                    )
            ray_indices = nerfacc.unpack_to_ray_indices(packed_info).long()

        zeros = torch.zeros_like(origins[:, :1])
        ray_samples = RaySamples(
            frustums=Frustums(
                origins=origins,
                directions=dirs,
                starts=starts,
                ends=ends,
                pixel_area=zeros,
            ),
            camera_indices=camera_indices,
        )
        return ray_samples, packed_info, ray_indices