# 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, List, Optional, Tuple

import nerfacc
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: Optional[int] = None,
        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,
        spacing_fn: Callable,
        spacing_fn_inv: Callable,
        num_samples: Optional[int] = None,
        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: Optional[int] = None,
        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: Optional[int] = None,
        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: Optional[int] = None,
        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: Optional[int] = None,
        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.
        include_original: Add original samples to ray.
        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.
    """

    def __init__(
        self,
        num_samples: Optional[int] = None,
        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.

        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.

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

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

        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,
        )

        num_samples = starts.shape[0]
        if num_samples == 0:
            # create a single fake sample and update packed_info accordingly
            # this says the last ray in packed_info has 1 sample, which starts and ends at 1
            packed_info[-1, 1] = 1
            starts = torch.ones((1, 1), dtype=starts.dtype, device=rays_o.device)
            ends = torch.ones((1, 1), dtype=ends.dtype, device=rays_o.device)

        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 and self.training:
            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


class ProposalNetworkSampler(Sampler):
    """Sampler that uses a proposal network to generate samples."""

    def __init__(
        self,
        num_proposal_samples_per_ray: int = 64,
        num_nerf_samples_per_ray: int = 32,
        num_proposal_network_iterations: int = 2,
        density_field: Optional[DensityGrid] = None,
        weight_threshold: float = 1e-2,
    ) -> None:
        super().__init__(density_field=density_field, weight_threshold=weight_threshold)
        self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
        self.num_nerf_samples_per_ray = num_nerf_samples_per_ray
        self.num_proposal_network_iterations = num_proposal_network_iterations

        # NOTE: a linear disparity sampler won't work until we sample in s space instead of t space
        self.uniform_sampler = UniformSampler()
        self.pdf_sampler = PDFSampler(include_original=False)

    def generate_ray_samples(
        self,
        ray_bundle: Optional[RayBundle] = None,
        density_fn: Optional[Callable] = None,
    ) -> Tuple[RaySamples, List, List]:

        weights_list = []
        ray_samples_list = []

        n = self.num_proposal_network_iterations
        for i_level in range(n + 1):
            is_prop = i_level < n
            num_samples = self.num_proposal_samples_per_ray if is_prop else self.num_nerf_samples_per_ray
            if i_level == 0:
                # Uniform sampling because we need to start with some samples
                ray_samples = self.uniform_sampler(ray_bundle, num_samples=num_samples)
            else:
                # PDF sampling based on the last samples and their weights
                ray_samples = self.pdf_sampler(ray_bundle, ray_samples, weights, num_samples=num_samples)
            if is_prop:
                density = density_fn(ray_samples.frustums.get_positions())
                weights = ray_samples.get_weights(density)
                weights_list.append(weights)  # (num_rays, num_samples)
                ray_samples_list.append(ray_samples)

        return ray_samples, weights_list, ray_samples_list