# Copyright (C) 2023, Gaussian-Grouping
# Gaussian-Grouping research group, https://github.com/lkeab/gaussian-grouping
# All rights reserved.
#
# ------------------------------------------------------------------------
# Modified from codes in Gaussian-Splatting 
# GRAPHDECO research group, https://team.inria.fr/graphdeco

from typing import NamedTuple
import torch.nn as nn
import torch
from . import _C

def cpu_deep_copy_tuple(input_tuple):
    copied_tensors = [item.cpu().clone() if isinstance(item, torch.Tensor) else item for item in input_tuple]
    return tuple(copied_tensors)

def rasterize_gaussians(
    prob,
    raster_settings
):
    return _RasterizeGaussians.apply(
        prob,
        raster_settings
    )

class _RasterizeGaussians(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        prob,
        raster_settings
    ):
        prob = torch.clamp(prob, min=0.0, max=1.0)
        # Restructure arguments the way that the C++ lib expects them
        args = (
            raster_settings.bg, 
            raster_settings.means3D,
            raster_settings.opacities,
            raster_settings.cov3Ds_precomp,
            raster_settings.viewmatrix,
            raster_settings.projmatrix,
            raster_settings.tanfovx,
            raster_settings.tanfovy,
            raster_settings.image_height,
            raster_settings.image_width,
            prob,
            raster_settings.prefiltered,
            raster_settings.debug
        )

        # Invoke C++/CUDA rasterizer
        if raster_settings.debug:
            cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted
            try:
                num_rendered, out_brightness, radii, geomBuffer, binningBuffer, imgBuffer = _C.rasterize_gaussians(*args)

            except Exception as ex:
                torch.save(cpu_args, "snapshot_fw.dump")
                print("\nAn error occured in forward. Please forward snapshot_fw.dump for debugging.")
                raise ex
        else:
            num_rendered, out_brightness, radii, geomBuffer, binningBuffer, imgBuffer = _C.rasterize_gaussians(*args)

        # Keep relevant tensors for backward
        ctx.raster_settings = raster_settings
        ctx.num_rendered = num_rendered
        ctx.save_for_backward(radii, prob, geomBuffer, binningBuffer, imgBuffer)
        return out_brightness

    @staticmethod
    def backward(ctx, grad_out_brightness):
        # Restore necessary values from context
        num_rendered = ctx.num_rendered
        raster_settings = ctx.raster_settings
        radii, prob, geomBuffer, binningBuffer, imgBuffer = ctx.saved_tensors

        # Restructure args as C++ method expects them
        args = (raster_settings.bg, # torch.Size([3])
                raster_settings.means3D,
                radii,  
                raster_settings.cov3Ds_precomp,
                raster_settings.viewmatrix, 
                raster_settings.projmatrix, 
                raster_settings.tanfovx, 
                raster_settings.tanfovy, 
                grad_out_brightness, # torch.Size([1, 504, 378])
                prob, # torch.Size([1418554, 1])
                geomBuffer, # torch.Size([112072936])
                num_rendered, # a <class 'int'> with value:  1572432
                binningBuffer, # torch.Size([57189247])
                imgBuffer, #  torch.Size([3048448])
                raster_settings.debug)

        # Compute gradients for relevant tensors by invoking backward method
        if raster_settings.debug:
            cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted
            try:
                grad_prob = _C.rasterize_gaussians_backward(*args)
            except Exception as ex:
                torch.save(cpu_args, "snapshot_bw.dump")
                print("\nAn error occured in backward. Writing snapshot_bw.dump for debugging.\n")
                raise ex
        else:
            grad_prob = _C.rasterize_gaussians_backward(*args)

        return grad_prob, None

class GaussianRasterizationSettings(NamedTuple):
    image_height: int
    image_width: int 
    tanfovx : float
    tanfovy : float
    bg : torch.Tensor
    viewmatrix : torch.Tensor
    projmatrix : torch.Tensor
    prefiltered : bool
    means3D: torch.Tensor
    opacities: torch.Tensor
    cov3Ds_precomp: torch.Tensor
    debug : bool

class GaussianRasterizer(nn.Module):
    def __init__(self, raster_settings):
        super().__init__()
        self.raster_settings = raster_settings

    def markVisible(self, positions):
        # Mark visible points (based on frustum culling for camera) with a boolean 
        with torch.no_grad():
            raster_settings = self.raster_settings
            visible = _C.mark_visible(
                positions,
                raster_settings.viewmatrix,
                raster_settings.projmatrix)
            
        return visible
    
    def forward(self, prob = None):
        
        raster_settings = self.raster_settings
        

        if prob is None:
            prob = torch.Tensor([])
        # Invoke C++/CUDA rasterization routine
        return rasterize_gaussians(
            prob,
            raster_settings
        )

