import torch
import kaolin
import math
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer


class CameraModule():
    def __init__(self):
        self.bg_color = torch.tensor([1.0] * 32).float()
        self.scale_modifier = 1.0

    def perspective_camera(self, points, camera_proj):
        projected_points = torch.bmm(points, camera_proj.permute(0, 2, 1))
        projected_2d_points = projected_points[:, :, :2] / projected_points[:, :, 2:3]

        return projected_2d_points

    def prepare_vertices(self, vertices, faces, camera_proj, camera_rot=None, camera_trans=None,
                     camera_transform=None):
        if camera_transform is None:
            assert camera_trans is not None and camera_rot is not None, \
                "camera_transform or camera_trans and camera_rot must be defined"
            vertices_camera = kaolin.render.camera.rotate_translate_points(vertices, camera_rot,
                                                            camera_trans)
        else:
            assert camera_trans is None and camera_rot is None, \
                "camera_trans and camera_rot must be None when camera_transform is defined"
            padded_vertices = torch.nn.functional.pad(
                vertices, (0, 1), mode='constant', value=1.
            )
            vertices_camera = (padded_vertices @ camera_transform)
        # Project the vertices on the camera image plan
        vertices_image = self.perspective_camera(vertices_camera, camera_proj)
        face_vertices_camera = kaolin.ops.mesh.index_vertices_by_faces(vertices_camera, faces)
        face_vertices_image = kaolin.ops.mesh.index_vertices_by_faces(vertices_image, faces)
        face_normals = kaolin.ops.mesh.face_normals(face_vertices_camera, unit=True)
        return face_vertices_camera, face_vertices_image, face_normals
    
    def render(self, data, resolution):
        verts_list = data['verts_list']
        faces_list = data['faces_list']
        verts_color_list = data['verts_color_list']

        B = len(verts_list)

        render_images = []
        render_soft_masks = []
        render_depths = []
        render_normals = []
        face_normals_list = []
        for b in range(B):
            intrinsics = data['intrinsics'][b]
            extrinsics = data['extrinsics'][b]
            #camera_proj = torch.stack([intrinsics[:, 0, 0] / intrinsics[:, 0, 2], intrinsics[:, 1, 1] / intrinsics[:, 1, 2], torch.ones_like(intrinsics[:, 0, 0])], -1).to(device)
            camera_proj = intrinsics
            camera_transform = extrinsics.permute(0, 2, 1)

            verts = verts_list[b].unsqueeze(0).repeat(intrinsics.shape[0], 1, 1)
            faces = faces_list[b]
            verts_color = verts_color_list[b].unsqueeze(0).repeat(intrinsics.shape[0], 1, 1)
            faces_color = verts_color[:, faces]

            face_vertices_camera, face_vertices_image, face_normals = self.prepare_vertices(
                verts, faces, camera_proj, camera_transform=camera_transform
            )
            face_vertices_image[:, :, :, 1] = -face_vertices_image[:, :, :, 1]
            #face_vertices_camera[:, :, :, 1:] = -face_vertices_camera[:, :, :, 1:]
            face_normals[:, :, 1:] = -face_normals[:, :, 1:]
            ### Perform Rasterization ###
            # Construct attributes that DI1-R rasterizer will interpolate.
            # the first is the UVS associated to each face
            # the second will make a hard segmentation mask
            face_attributes = [
                faces_color,
                torch.ones((faces_color.shape[0], faces_color.shape[1], 3, 1), device=verts.device),
                face_vertices_camera[:, :, :, 2:],
                face_normals.unsqueeze(-2).repeat(1, 1, 3, 1),
            ]

            # If you have nvdiffrast installed you can change rast_backend to
            # nvdiffrast or nvdiffrast_fwd
            image_features, soft_masks, face_idx = kaolin.render.mesh.dibr_rasterization(
                resolution, resolution, -face_vertices_camera[:, :, :, -1],
                face_vertices_image, face_attributes, face_normals[:, :, -1],
                rast_backend='cuda')

            # image_features is a tuple in composed of the interpolated attributes of face_attributes
            images, masks, depths, normals = image_features
            images = torch.clamp(images * masks, 0., 1.)
            depths = (depths * masks)
            normals = (normals * masks)
            
            render_images.append(images)
            render_soft_masks.append(soft_masks)
            render_depths.append(depths)
            render_normals.append(normals)
            face_normals_list.append(face_normals)

        render_images = torch.stack(render_images, 0)
        render_soft_masks = torch.stack(render_soft_masks, 0)
        render_depths = torch.stack(render_depths, 0)
        render_normals = torch.stack(render_normals, 0)

        data['render_images'] = render_images
        data['render_soft_masks'] = render_soft_masks
        data['render_depths'] = render_depths
        data['render_normals'] = render_normals
        data['verts_list'] = verts_list
        data['faces_list'] = faces_list
        data['face_normals_list'] = face_normals_list

        return data
    

    def render_gaussian(self, data, resolution):
        """
        Render the scene. 
        
        Background tensor (bg_color) must be on GPU!
        """
        B = data['xyz'].shape[0]
        xyz = data['xyz']
        #shs = rearrange(data['shs'], 'b n (x y) -> b n x y', y=3)
        colors_precomp = data['color']
        opacity = data['opacity']
        scales = data['scales']
        rotations = data['rotation']
        fovx = data['fovx']
        fovy = data['fovy']
        bg_color = self.bg_color if 'bg_color' not in data.keys() else data['bg_color']

        world_view_transform = data['world_view_transform']
        full_proj_transform = data['full_proj_transform']
        camera_center = data['camera_center']
    
        # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
        screenspace_points = torch.zeros_like(xyz, dtype=xyz.dtype, requires_grad=True, device=xyz.device) + 0
        try:
            screenspace_points.retain_grad()
        except:
            pass

        render_images = []
        radii = []
        for b in range(B):

            tanfovx = math.tan(fovx[b] * 0.5)
            tanfovy = math.tan(fovy[b] * 0.5)
            # Set up rasterization configuration
            raster_settings = GaussianRasterizationSettings(
                image_height=int(resolution),
                image_width=int(resolution),
                tanfovx=tanfovx,
                tanfovy=tanfovy,
                bg=bg_color.to(xyz.device),
                scale_modifier=self.scale_modifier,
                viewmatrix=world_view_transform[b],
                projmatrix=full_proj_transform[b],
                sh_degree=0,
                campos=camera_center[b],
                prefiltered=False,
                debug=False
            )

            rasterizer = GaussianRasterizer(raster_settings=raster_settings)

            means3D = xyz[b]
            means2D = screenspace_points[b]
            
            # Rasterize visible Gaussians to image, obtain their radii (on screen). 
            render_images_b, radii_b, _ = rasterizer(
                means3D = means3D,
                means2D = means2D,
                #shs = shs[b],
                colors_precomp = colors_precomp[b],
                opacities = opacity[b],
                scales = scales[b],
                rotations = rotations[b])
            render_images.append(render_images_b)
            radii.append(radii_b)

        render_images = torch.stack(render_images)
        radii = torch.stack(radii)
        data['render_images'] = render_images
        data['viewspace_points'] =  screenspace_points
        data['visibility_filter'] = radii > 0
        data['radii'] = radii
        return data