import math
from typing import Callable, List, Optional, Tuple, Union

import torch
from ..grid import _enlarge_aabb
from ..volrend import (
    render_visibility_from_alpha,
    render_visibility_from_density,
)
from .base import AbstractEstimator
from torch import Tensor

try:
    import svox
except ImportError:
    raise ImportError(
        "Please install this forked version of svox: "
        "pip install git+https://github.com/liruilong940607/svox.git"
    )


class N3TreeEstimator(AbstractEstimator):
    """Use N3Tree to implement Occupancy Grid.
    
    This allows more flexible topologies than the cascaded grid. However, it is
    slower to create samples from the tree than the cascaded grid. By default,
    it has the same topology as the cascaded grid but `self.tree` can be
    modified to have different topologies.
    """

    def __init__(
        self,
        roi_aabb: Union[List[int], Tensor],
        resolution: Union[int, List[int], Tensor] = 128,
        levels: int = 1,
        **kwargs,
    ) -> None:
        super().__init__()

        if "contraction_type" in kwargs:
            raise ValueError(
                "`contraction_type` is not supported anymore for nerfacc >= 0.4.0."
            )

        # check the resolution is legal
        assert isinstance(resolution, int), "N3Tree only supports uniform resolution!"

        # check the roi_aabb is legal
        if isinstance(roi_aabb, (list, tuple)):
            roi_aabb = torch.tensor(roi_aabb, dtype=torch.float32)
        assert isinstance(roi_aabb, Tensor), f"Invalid type: {roi_aabb}!"
        assert roi_aabb.shape[0] == 6, f"Invalid shape: {roi_aabb}!"
        roi_aabb = roi_aabb.cpu()

        # to be compatible with the OccupancyGrid
        aabbs = torch.stack(
            [_enlarge_aabb(roi_aabb, 2**i) for i in range(levels)], dim=0
        )
        self.register_buffer("aabbs", aabbs)  # [n_aabbs, 6]

        center = (roi_aabb[:3] + roi_aabb[3:]) / 2.0
        radius = (roi_aabb[3:] - roi_aabb[:3]) / 2.0 * 2 ** (levels - 1)

        tree_depth = int(math.log2(resolution)) - 1
        self.tree = svox.N3Tree(
            N=2,
            data_dim=1,
            init_refine=tree_depth,
            depth_limit=20,
            radius=radius.tolist(),
            center=center.tolist(),
        )
        _aabbs = [_enlarge_aabb(roi_aabb, 2**i) for i in range(levels - 1)]
        for aabb in _aabbs[::-1]:
            leaf_c = self.tree.corners + self.tree.lengths * 0.5
            sel = ((leaf_c > aabb[:3]) & (leaf_c < aabb[3:])).all(dim=-1)
            self.tree[sel].refine()
        # print("tree size", len(self.tree), "at resolution", resolution)

        self.thresh = 0.0

    @torch.no_grad()
    def sampling(
        self,
        # rays
        rays_o: Tensor,  # [n_rays, 3]
        rays_d: Tensor,  # [n_rays, 3]
        # sigma/alpha function for skipping invisible space
        sigma_fn: Optional[Callable] = None,
        alpha_fn: Optional[Callable] = None,
        near_plane: float = 0.0,
        far_plane: float = 1e10,
        t_min: Optional[Tensor] = None,  # [n_rays]
        t_max: Optional[Tensor] = None,  # [n_rays]
        # rendering options
        render_step_size: float = 1e-3,
        early_stop_eps: float = 1e-4,
        alpha_thre: float = 0.0,
        stratified: bool = False,
        cone_angle: float = 0.0,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """Sampling with spatial skipping.

        Note:
            This function is not differentiable to any inputs.

        Args:
            rays_o: Ray origins of shape (n_rays, 3).
            rays_d: Normalized ray directions of shape (n_rays, 3).
            sigma_fn: Optional. If provided, the marching will skip the invisible space
                by evaluating the density along the ray with `sigma_fn`. It should be a
                function that takes in samples {t_starts (N,), t_ends (N,),
                ray indices (N,)} and returns the post-activation density values (N,).
                You should only provide either `sigma_fn` or `alpha_fn`.
            alpha_fn: Optional. If provided, the marching will skip the invisible space
                by evaluating the density along the ray with `alpha_fn`. It should be a
                function that takes in samples {t_starts (N,), t_ends (N,),
                ray indices (N,)} and returns the post-activation opacity values (N,).
                You should only provide either `sigma_fn` or `alpha_fn`.
            near_plane: Optional. Near plane distance. Default: 0.0.
            far_plane: Optional. Far plane distance. Default: 1e10.
            t_min: Optional. Per-ray minimum distance. Tensor with shape (n_rays).
                If profided, the marching will start from maximum of t_min and near_plane.
            t_max: Optional. Per-ray maximum distance. Tensor with shape (n_rays).
                If profided, the marching will stop by minimum of t_max and far_plane.
            render_step_size: Step size for marching. Default: 1e-3.
            early_stop_eps: Early stop threshold for skipping invisible space. Default: 1e-4.
            alpha_thre: Alpha threshold for skipping empty space. Default: 0.0.
            stratified: Whether to use stratified sampling. Default: False.
            cone_angle: Cone angle for linearly-increased step size. 0. means
                constant step size. Default: 0.0.

        Returns:
            A tuple of {LongTensor, Tensor, Tensor}:

            - **ray_indices**: Ray index of each sample. IntTensor with shape (n_samples).
            - **t_starts**: Per-sample start distance. Tensor with shape (n_samples,).
            - **t_ends**: Per-sample end distance. Tensor with shape (n_samples,).

        Examples:

        .. code-block:: python

            >>> ray_indices, t_starts, t_ends = grid.sampling(
            >>>     rays_o, rays_d, render_step_size=1e-3)
            >>> t_mid = (t_starts + t_ends) / 2.0
            >>> sample_locs = rays_o[ray_indices] + t_mid * rays_d[ray_indices]

        """

        assert t_min is None and t_max is None, (
            "Do not supported per-ray min max. Please use near_plane and far_plane instead."
        )
        if stratified:
            near_plane += torch.rand(()).item() * render_step_size

        t_starts, t_ends, packed_info, ray_indices = svox.volume_sample(
            self.tree,
            thresh=self.thresh,
            rays=svox.Rays(rays_o.contiguous(), rays_d.contiguous(), rays_d.contiguous()),
            step_size=render_step_size,
            cone_angle=cone_angle,
            near_plane=near_plane,
            far_plane=far_plane,
        )
        packed_info = packed_info.long()
        ray_indices = ray_indices.long()

        # skip invisible space
        if (alpha_thre > 0.0 or early_stop_eps > 0.0) and (
            sigma_fn is not None or alpha_fn is not None
        ):
            alpha_thre = min(alpha_thre, self.thresh)

            # Compute visibility of the samples, and filter out invisible samples
            if sigma_fn is not None:
                if t_starts.shape[0] != 0:
                    sigmas = sigma_fn(t_starts, t_ends, ray_indices)
                else:
                    sigmas = torch.empty((0,), device=t_starts.device)
                assert (
                    sigmas.shape == t_starts.shape
                ), "sigmas must have shape of (N,)! Got {}".format(sigmas.shape)
                masks = render_visibility_from_density(
                    t_starts=t_starts,
                    t_ends=t_ends,
                    sigmas=sigmas,
                    ray_indices=ray_indices,
                    n_rays=len(rays_o),
                    early_stop_eps=early_stop_eps,
                    alpha_thre=alpha_thre,
                )
            elif alpha_fn is not None:
                if t_starts.shape[0] != 0:
                    alphas = alpha_fn(t_starts, t_ends, ray_indices)
                else:
                    alphas = torch.empty((0,), device=t_starts.device)
                assert (
                    alphas.shape == t_starts.shape
                ), "alphas must have shape of (N,)! Got {}".format(alphas.shape)
                masks = render_visibility_from_alpha(
                    alphas=alphas,
                    ray_indices=ray_indices,
                    n_rays=len(rays_o),
                    early_stop_eps=early_stop_eps,
                    alpha_thre=alpha_thre,
                )
            ray_indices, t_starts, t_ends = (
                ray_indices[masks],
                t_starts[masks],
                t_ends[masks],
            )
        return ray_indices, t_starts, t_ends

    @torch.no_grad()
    def update_every_n_steps(
        self,
        step: int,
        occ_eval_fn: Callable,
        occ_thre: float = 1e-2,
        ema_decay: float = 0.95,
        warmup_steps: int = 256,
        n: int = 16,
    ) -> None:
        """Update the estimator every n steps during training.

        Args:
            step: Current training step.
            occ_eval_fn: A function that takes in sample locations :math:`(N, 3)` and
                returns the occupancy values :math:`(N, 1)` at those locations.
            occ_thre: Threshold used to binarize the occupancy grid. Default: 1e-2.
            ema_decay: The decay rate for EMA updates. Default: 0.95.
            warmup_steps: Sample all cells during the warmup stage. After the warmup
                stage we change the sampling strategy to 1/4 uniformly sampled cells
                together with 1/4 occupied cells. Default: 256.
            n: Update the grid every n steps. Default: 16.
        """
        if not self.training:
            raise RuntimeError(
                "You should only call this function only during training. "
                "Please call _update() directly if you want to update the "
                "field during inference."
            )
        if step % n == 0 and self.training:
            self._update(
                step=step,
                occ_eval_fn=occ_eval_fn,
                occ_thre=occ_thre,
                ema_decay=ema_decay,
                warmup_steps=warmup_steps,
            )

    @torch.no_grad()
    def _sample_uniform_and_occupied_cells(self, n: int) -> List[Tensor]:
        """Samples both n uniform and occupied cells."""
        uniform_indices = torch.randint(len(self.tree), (n,), device=self.device)
        occupied_indices = torch.nonzero(self.tree[:].values >= self.thresh)[:, 0]
        if n < len(occupied_indices):
            selector = torch.randint(len(occupied_indices), (n,), device=self.device)
            occupied_indices = occupied_indices[selector]
        indices = torch.cat([uniform_indices, occupied_indices], dim=0)
        return indices

    @torch.no_grad()
    def _update(
        self,
        step: int,
        occ_eval_fn: Callable,
        occ_thre: float = 0.01,
        ema_decay: float = 0.95,
        warmup_steps: int = 256,
    ) -> None:
        """Update the occ field in the EMA way."""
        if step < warmup_steps:
            x = self.tree.sample(1).squeeze(1)
            occ = occ_eval_fn(x).squeeze(-1)
            sel = (*self.tree._all_leaves().T,)
            self.tree.data.data[sel] = torch.maximum(self.tree.data.data[sel] * ema_decay, occ[:, None])
        else:
            N = len(self.tree) // 4
            indices = self._sample_uniform_and_occupied_cells(N)
            x = self.tree[indices].sample(1).squeeze(1)
            occ = occ_eval_fn(x).squeeze(-1)
            self.tree[indices] = torch.maximum(
                self.tree[indices].values * ema_decay, occ[:, None]
            )
        self.thresh = min(occ_thre, self.tree[:].values.mean().item())


if __name__ == "__main__":
    roi_aabb = [-1.0, -1.0, -1.0, 1.0, 1.0, 1.0]
    resolution = 128
    levels = 4
    estimator = N3TreeEstimator(roi_aabb, resolution, levels)

    def occ_eval_fn(x):
        return torch.rand(len(x), 1)

    estimator.update_every_n_steps(0, occ_eval_fn, occ_thre=0.5)
