# ORIGINAL LICENSE
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Modified by Zexin He
# The modifications are subject to the same license as the original.


import itertools
import torch
import torch.nn as nn
import torch.nn.functional as F

from .utils.renderer import ImportanceRenderer, sample_from_planes
from .utils.ray_sampler import RaySampler
from ...utils.ops import get_rank


class OSGDecoder(nn.Module):
    """
    Triplane decoder that gives RGB and sigma values from sampled features.
    Using ReLU here instead of Softplus in the original implementation.
    
    Reference:
    EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L112
    """
    def __init__(self, n_features: int,
                 hidden_dim: int = 64, 
                 num_layers: int = 2, 
                 activation: nn.Module = nn.ReLU,
                 sdf_bias='sphere',
                 sdf_bias_params=0.5,
                 output_normal=True,
                 normal_type='finite_difference'):
        super().__init__()
        self.sdf_bias = sdf_bias
        self.sdf_bias_params = sdf_bias_params
        self.output_normal = output_normal
        self.normal_type = normal_type
        self.net = nn.Sequential(
            nn.Linear(3 * n_features, hidden_dim),
            activation(),
            *itertools.chain(*[[
                nn.Linear(hidden_dim, hidden_dim),
                activation(),
            ] for _ in range(num_layers - 2)]),
            nn.Linear(hidden_dim, 1 + 3),
        )
        # init all bias to zero
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.zeros_(m.bias)

    def forward(self, ray_directions, sample_coordinates, plane_axes, planes, options):
        # Aggregate features by mean
        # sampled_features = sampled_features.mean(1)
        # Aggregate features by concatenation
        # torch.set_grad_enabled(True)
        # sample_coordinates.requires_grad_(True)

        sampled_features = sample_from_planes(plane_axes, planes, sample_coordinates, padding_mode='zeros', box_warp=options['box_warp'])

        
        _N, n_planes, _M, _C = sampled_features.shape
        sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
        x = sampled_features

        N, M, C = x.shape
        # x = x.contiguous().view(N*M, C)

        x = self.net(x)
        x = x.view(N, M, -1)
        rgb = torch.sigmoid(x[..., 1:])*(1 + 2*0.001) - 0.001  # Uses sigmoid clamping from MipNeRF
        
        sdf = x[..., 0:1]
        # import ipdb; ipdb.set_trace()
        # print(f'sample_coordinates shape: {sample_coordinates.shape}')
        # sdf = self.get_shifted_sdf(sample_coordinates, sdf)

        # calculate normal
        eps = 0.01
        offsets = torch.as_tensor(
            [[eps, 0.0, 0.0], [0.0, eps, 0.0], [0.0, 0.0, eps]]
        ).to(sample_coordinates)
        points_offset = (
            sample_coordinates[..., None, :] + offsets # Float[Tensor, "... 3 3"]
        ).clamp(options['sampler_bbox_min'], options['sampler_bbox_max'])

        sdf_offset_list = [self.forward_sdf(
            plane_axes,
            planes,
            points_offset[:,:,i,:],
            options
        ).unsqueeze(-2) for i in range(points_offset.shape[-2])] # Float[Tensor, "... 3 1"]
        # import ipdb; ipdb.set_trace()
        
        sdf_offset = torch.cat(sdf_offset_list, -2)
        sdf_grad = (sdf_offset[..., 0::1, 0] - sdf) / eps
        
        normal = F.normalize(sdf_grad, dim=-1).to(sdf.dtype)
        return {'rgb': rgb, 'sdf': sdf, 'normal': normal, 'sdf_grad': sdf_grad}
    
    def forward_sdf(self, plane_axes, planes, points_offset, options):

        sampled_features = sample_from_planes(plane_axes, planes, points_offset, padding_mode='zeros', box_warp=options['box_warp'])
        _N, n_planes, _M, _C = sampled_features.shape
        sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
        x = sampled_features

        N, M, C = x.shape
        # x = x.contiguous().view(N*M, C)

        x = self.net(x)
        x = x.view(N, M, -1)
        sdf = x[..., 0:1]
        # sdf = self.get_shifted_sdf(points_offset, sdf)
        return sdf
    
    def get_shifted_sdf(
        self, points, sdf
    ):
        if self.sdf_bias == "sphere":
            assert isinstance(self.sdf_bias_params, float)
            radius = self.sdf_bias_params
            sdf_bias = (points**2).sum(dim=-1, keepdim=True).sqrt() - radius
        else:
            raise ValueError(f"Unknown sdf bias {self.cfg.sdf_bias}")
        return sdf + sdf_bias.to(sdf.dtype)
    
    
class TriplaneSynthesizer(nn.Module):
    """
    Synthesizer that renders a triplane volume with planes and a camera.
    
    Reference:
    EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L19
    """

    DEFAULT_RENDERING_KWARGS = {
        'ray_start': 'auto',
        'ray_end': 'auto',
        'box_warp': 1.2,
        # 'box_warp': 1.,
        'white_back': True,
        'disparity_space_sampling': False,
        'clamp_mode': 'softplus',
        # 'sampler_bbox_min': -1,
        # 'sampler_bbox_max': 1.,
        'sampler_bbox_min': -0.6,
        'sampler_bbox_max': 0.6,
    }
    print('DEFAULT_RENDERING_KWARGS')
    print(DEFAULT_RENDERING_KWARGS)


    def __init__(self, triplane_dim: int, samples_per_ray: int, osg_decoder='default'):
        super().__init__()

        # attributes
        self.triplane_dim = triplane_dim
        self.rendering_kwargs = {
            **self.DEFAULT_RENDERING_KWARGS,
            'depth_resolution': samples_per_ray,
            'depth_resolution_importance': 0
            # 'depth_resolution': samples_per_ray // 2,
            # 'depth_resolution_importance': samples_per_ray // 2,
        }

        # renderings
        self.renderer = ImportanceRenderer()
        self.ray_sampler = RaySampler()
        # modules
        if osg_decoder == 'default':
            self.decoder = OSGDecoder(n_features=triplane_dim)
        else:
            raise NotImplementedError

    def forward(self, planes, ray_origins, ray_directions, render_size, bgcolor=None):
        # planes: (N, 3, D', H', W')
        # render_size: int
        assert ray_origins.dim() == 3, "ray_origins should be 3-dimensional"
        

        # Perform volume rendering
        rgb_samples, depth_samples, weights_samples, sdf_grad, normal_samples = self.renderer(
            planes, self.decoder, ray_origins, ray_directions, self.rendering_kwargs, bgcolor
        )
        N = planes.shape[0]

        # zhaohx : add for normals
        normal_samples = F.normalize(normal_samples, dim=-1)
        normal_samples = (normal_samples + 1.0) / 2.0   # for visualization
        normal_samples = torch.lerp(torch.zeros_like(normal_samples), normal_samples, weights_samples)

        # Reshape into 'raw' neural-rendered image
        Himg = Wimg = render_size
        rgb_images = rgb_samples.permute(0, 2, 1).reshape(N, rgb_samples.shape[-1], Himg, Wimg).contiguous()
        depth_images = depth_samples.permute(0, 2, 1).reshape(N, 1, Himg, Wimg)
        weight_images = weights_samples.permute(0, 2, 1).reshape(N, 1, Himg, Wimg)

        # zhaohx : add for normals
        normal_images = normal_samples.permute(0, 2, 1).reshape(N, normal_samples.shape[-1], Himg, Wimg).contiguous()

        # return {
        #     'images_rgb': rgb_images,
        #     'images_depth': depth_images,
        #     'images_weight': weight_images,
        # }

        return {
            'comp_rgb': rgb_images,
            'comp_depth': depth_images,
            'opacity': weight_images,
            'sdf_grad': sdf_grad,
            'comp_normal': normal_images
        }
        # 输出normal的话在这个return里加

    def forward_grid(self, planes, grid_size: int, aabb: torch.Tensor = None):
        # planes: (N, 3, D', H', W')
        # grid_size: int
        # aabb: (N, 2, 3)
        if aabb is None:
            aabb = torch.tensor([
                [self.rendering_kwargs['sampler_bbox_min']] * 3,
                [self.rendering_kwargs['sampler_bbox_max']] * 3,
            ], device=planes.device, dtype=planes.dtype).unsqueeze(0).repeat(planes.shape[0], 1, 1)
        assert planes.shape[0] == aabb.shape[0], "Batch size mismatch for planes and aabb"
        N = planes.shape[0]

        # create grid points for triplane query
        grid_points = []
        for i in range(N):
            grid_points.append(torch.stack(torch.meshgrid(
                torch.linspace(aabb[i, 0, 0], aabb[i, 1, 0], grid_size, device=planes.device),
                torch.linspace(aabb[i, 0, 1], aabb[i, 1, 1], grid_size, device=planes.device),
                torch.linspace(aabb[i, 0, 2], aabb[i, 1, 2], grid_size, device=planes.device),
                indexing='ij',
            ), dim=-1).reshape(-1, 3))
        cube_grid = torch.stack(grid_points, dim=0).to(planes.device)

        features = self.forward_points(planes, cube_grid)

        # reshape into grid
        features = {
            k: v.reshape(N, grid_size, grid_size, grid_size, -1)
            for k, v in features.items()
        }
        return features

    def forward_points(self, planes, points: torch.Tensor, chunk_size: int = 2**20):
        # planes: (N, 3, D', H', W')
        # points: (N, P, 3)
        N, P = points.shape[:2]

        # query triplane in chunks
        outs = []
        for i in range(0, points.shape[1], chunk_size):
            chunk_points = points[:, i:i+chunk_size]

            # query triplane
            # chunk_out = self.renderer.run_model_activated(
            chunk_out = self.renderer.run_model(
                planes=planes,
                decoder=self.decoder,
                sample_coordinates=chunk_points,
                sample_directions=torch.zeros_like(chunk_points),
                options=self.rendering_kwargs,
            )
            outs.append(chunk_out)

        # concatenate the outputs
        point_features = {
            k: torch.cat([out[k] for out in outs], dim=1)
            for k in outs[0].keys()
        }
        return point_features