# part of the code from 
#   https://github.com/benjiebob/SMALify/blob/master/smal_fitter/p3d_renderer.py

import torch
import torch.nn.functional as F
from scipy.io  import loadmat
import numpy as np
# import config

import pytorch3d
from pytorch3d.structures import Meshes
from pytorch3d.renderer import (
    PerspectiveCameras, look_at_view_transform, look_at_rotation, 
    RasterizationSettings, MeshRenderer, MeshRasterizer, BlendParams,
    PointLights, HardPhongShader, SoftSilhouetteShader, Materials, Textures, 
    DirectionalLights
)
from pytorch3d.renderer import TexturesVertex, SoftPhongShader
from pytorch3d.io import load_objs_as_meshes

MESH_COLOR_0 = [0, 172, 223]
MESH_COLOR_1 = [172, 223, 0]


'''
Explanation of the shift between projection results from opendr and pytorch3d:
    (0, 0, ?) will be projected to 127.5 (pytorch3d) instead of 128 (opendr)
    imagine you have an image of size 4:
    middle of the first pixel is 0
    middle of the last pixel is 3
    => middle of the imgae would be 1.5 and not 2!
    so in order to go from pytorch3d predictions to opendr we would calculate: p_odr = p_p3d * (128/127.5)
To reproject points (p3d) by hand according to this pytorch3d renderer we would do the following steps:
    1.) build camera matrix
        K = np.array([[flength, 0, c_x],
                    [0, flength, c_y],
                    [0, 0, 1]], np.float)
    2.) we don't need to add extrinsics, as the mesh comes with translation (which is 
        added within smal_pytorch). all 3d points are already in the camera coordinate system.
        -> projection reduces to p2d_proj = K*p3d
    3.) convert to pytorch3d conventions (0 in the middle of the first pixel)
        p2d_proj_pytorch3d = p2d_proj / image_size * (image_size-1.)
renderer.py - project_points_p3d: shows an example of what is described above, but 
    same focal length for the whole batch

'''

class SilhRenderer(torch.nn.Module):
    def __init__(self, image_size, adapt_R_wldo=False):
        super(SilhRenderer, self).__init__()
        # see: https://pytorch3d.org/files/fit_textured_mesh.py, line 315
        # adapt_R=True is True for all my experiments
        # image_size: one number, integer
        # -----
        # set mesh color
        self.register_buffer('mesh_color_0', torch.FloatTensor(MESH_COLOR_0))
        self.register_buffer('mesh_color_1', torch.FloatTensor(MESH_COLOR_1))
        # prepare extrinsics, which in our case don't change
        R = torch.Tensor(np.eye(3)).float()[None, :, :]
        T = torch.Tensor(np.zeros((1, 3))).float()
        if adapt_R_wldo:
            R[0, 0, 0] = -1
        else:       # used for all my own experiments
            R[0, 0, 0] = -1
            R[0, 1, 1] = -1            
        self.register_buffer('R', R)
        self.register_buffer('T', T)
        # prepare that part of the intrinsics which does not change either
        # principal_point_prep = torch.Tensor([self.image_size / 2., self.image_size / 2.]).float()[None, :].float().to(device) 
        # image_size_prep = torch.Tensor([self.image_size, self.image_size]).float()[None, :].float().to(device) 
        self.img_size_scalar = image_size
        self.register_buffer('image_size', torch.Tensor([image_size, image_size]).float()[None, :].float())
        self.register_buffer('principal_point', torch.Tensor([image_size / 2., image_size / 2.]).float()[None, :].float())
        # Rasterization settings for differentiable rendering, where the blur_radius
        # initialization is based on Liu et al, 'Soft Rasterizer: A Differentiable 
        # Renderer for Image-based 3D Reasoning', ICCV 2019
        self.blend_params = BlendParams(sigma=1e-4, gamma=1e-4)
        self.raster_settings_soft = RasterizationSettings(
            image_size=image_size,      # 128
            blur_radius=np.log(1. / 1e-4 - 1.)*self.blend_params.sigma,  
            faces_per_pixel=100)     #50, 
        # Renderer for Image-based 3D Reasoning', body part segmentation
        self.blend_params_parts = BlendParams(sigma=2*1e-4, gamma=1e-4)
        self.raster_settings_soft_parts = RasterizationSettings(
            image_size=image_size,      # 128
            blur_radius=np.log(1. / 1e-4 - 1.)*self.blend_params_parts.sigma,  
            faces_per_pixel=60)     #50, 
        # settings for visualization renderer
        self.raster_settings_vis = RasterizationSettings(
                image_size=image_size, 
                blur_radius=0.0, 
                faces_per_pixel=1)

    def _get_cam(self, focal_lengths):
        device = focal_lengths.device
        bs = focal_lengths.shape[0]
        if pytorch3d.__version__ == '0.2.5':
            cameras = PerspectiveCameras(device=device,
                focal_length=focal_lengths.repeat((1, 2)), 
                principal_point=self.principal_point.repeat((bs, 1)), 
                R=self.R.repeat((bs, 1, 1)), T=self.T.repeat((bs, 1)), 
                image_size=self.image_size.repeat((bs, 1)))
        elif pytorch3d.__version__ == '0.6.1':
            cameras = PerspectiveCameras(device=device, in_ndc=False,
            focal_length=focal_lengths.repeat((1, 2)), 
            principal_point=self.principal_point.repeat((bs, 1)), 
            R=self.R.repeat((bs, 1, 1)), T=self.T.repeat((bs, 1)), 
            image_size=self.image_size.repeat((bs, 1)))
        else: 
            print('this part depends on the version of pytorch3d, code was developed with 0.2.5')
            raise ValueError
        return cameras

    def _get_visualization_from_mesh(self, mesh, cameras, lights=None):
        # color renderer for visualization
        with torch.no_grad():
            device = mesh.device
            # renderer for visualization
            if lights is None:
                lights = PointLights(device=device, location=[[0.0, 0.0, 3.0]])
            vis_renderer = MeshRenderer(
                rasterizer=MeshRasterizer(
                    cameras=cameras, 
                    raster_settings=self.raster_settings_vis),
                shader=HardPhongShader(
                    device=device, 
                    cameras=cameras,
                    lights=lights))
            # render image:
            visualization = vis_renderer(mesh).permute(0, 3, 1, 2)[:, :3, :, :]
        return visualization


    def calculate_vertex_visibility(self, vertices, faces, focal_lengths, soft=False):
        tex = torch.ones_like(vertices) * self.mesh_color_0    # (1, V, 3)
        textures = Textures(verts_rgb=tex)
        mesh = Meshes(verts=vertices, faces=faces, textures=textures)
        cameras = self._get_cam(focal_lengths)
        # NEW: use the rasterizer to check vertex visibility
        #   see: https://github.com/facebookresearch/pytorch3d/issues/126
        # Get a rasterizer
        if soft:
            rasterizer = MeshRasterizer(cameras=cameras, 
                                        raster_settings=self.raster_settings_soft)
        else:
            rasterizer = MeshRasterizer(cameras=cameras, 
                            raster_settings=self.raster_settings_vis)
        # Get the output from rasterization
        fragments = rasterizer(mesh)
        # pix_to_face is of shape (N, H, W, 1)
        pix_to_face = fragments.pix_to_face  
        # (F, 3) where F is the total number of faces across all the meshes in the batch
        packed_faces = mesh.faces_packed() 
        # (V, 3) where V is the total number of verts across all the meshes in the batch
        packed_verts = mesh.verts_packed() 
        vertex_visibility_map = torch.zeros(packed_verts.shape[0])   # (V,)
        # Indices of unique visible faces
        visible_faces = pix_to_face.unique()    # [0]   # (num_visible_faces )
        # Get Indices of unique visible verts using the vertex indices in the faces
        visible_verts_idx = packed_faces[visible_faces]    # (num_visible_faces,  3)
        unique_visible_verts_idx = torch.unique(visible_verts_idx)   # (num_visible_verts, )
        # Update visibility indicator to 1 for all visible vertices 
        vertex_visibility_map[unique_visible_verts_idx] = 1.0
        # since all meshes have the same amount of vertices, we can reshape the result 
        bs = vertices.shape[0]
        vertex_visibility_map_resh = vertex_visibility_map.reshape((bs, -1))
        return pix_to_face, vertex_visibility_map_resh


    def get_torch_meshes(self, vertices, faces, color=0):
        # create pytorch mesh
        if color == 0:
            mesh_color = self.mesh_color_0
        else:
            mesh_color = self.mesh_color_1
        tex = torch.ones_like(vertices) * mesh_color    # (1, V, 3)
        textures = Textures(verts_rgb=tex)
        mesh = Meshes(verts=vertices, faces=faces, textures=textures)  
        return mesh


    def get_visualization_nograd(self, vertices, faces, focal_lengths, color=0):
        # vertices: torch.Size([bs, 3889, 3])
        # faces: torch.Size([bs, 7774, 3]), int
        # focal_lengths: torch.Size([bs, 1])
        device = vertices.device
        # create cameras
        cameras = self._get_cam(focal_lengths)
        # create pytorch mesh
        if color == 0:
            mesh_color = self.mesh_color_0      # blue
        elif color == 1:
            mesh_color = self.mesh_color_1      
        elif color ==  2:   
            MESH_COLOR_2 = [240, 250, 240]      # white
            mesh_color =  torch.FloatTensor(MESH_COLOR_2).to(device)
        elif color ==  3:   
            # MESH_COLOR_3 = [223, 0, 172]         # pink
            # MESH_COLOR_3 = [245, 245, 220]         # beige
            MESH_COLOR_3 = [166, 173, 164] 
            mesh_color =  torch.FloatTensor(MESH_COLOR_3).to(device)        
        else:
            MESH_COLOR_2 = [240, 250, 240]
            mesh_color =  torch.FloatTensor(MESH_COLOR_2).to(device)
        tex = torch.ones_like(vertices) * mesh_color    # (1, V, 3)
        textures = Textures(verts_rgb=tex)
        mesh = Meshes(verts=vertices, faces=faces, textures=textures)  
        # render mesh (no gradients)
        # lights = PointLights(device=device, location=[[0.0, 0.0, 3.0]])
        # lights = PointLights(device=device, location=[[2.0, 2.0, -2.0]])
        lights = DirectionalLights(device=device, direction=[[0.0, -5.0, -10.0]])
        visualization = self._get_visualization_from_mesh(mesh, cameras, lights=lights)
        return visualization
              
    def project_points(self, points, focal_lengths=None, cameras=None):
        # points: torch.Size([bs, n_points, 3])
        # either focal_lengths or cameras is needed:
        #   focal_lenghts: torch.Size([bs, 1])
        #   cameras: pytorch camera, for example PerspectiveCameras()
        bs = points.shape[0]
        device = points.device
        screen_size = self.image_size.repeat((bs, 1))
        if cameras is None:
            cameras = self._get_cam(focal_lengths)
        if pytorch3d.__version__ == '0.2.5':
            proj_points_orig = cameras.transform_points_screen(points, screen_size)[:, :, [1, 0]]       # used in the original virtuel environment (for cvpr BARC submission) 
        elif pytorch3d.__version__ == '0.6.1':
            proj_points_orig = cameras.transform_points_screen(points)[:, :, [1, 0]]        
        else: 
            print('this part depends on the version of pytorch3d, code was developed with 0.2.5')
            raise ValueError
        # flip, otherwise the 1st and 2nd row are exchanged compared to the ground truth
        proj_points = torch.flip(proj_points_orig, [2])   
        # --- project points 'manually'
        # j_proj = project_points_p3d(image_size, focal_length, points, device)
        return proj_points      

    def forward(self, vertices, points, faces, focal_lengths, color=None):
        # vertices: torch.Size([bs, 3889, 3])
        # points: torch.Size([bs, n_points, 3]) (or None)
        # faces: torch.Size([bs, 7774, 3]), int
        # focal_lengths: torch.Size([bs, 1])
        # color: if None we don't render a visualization, else it should
        #   either be 0 or 1
        # ---> important: results are around 0.5 pixels off compared to chumpy!
        #   have a look at renderer.py for an explanation
        # create cameras
        cameras = self._get_cam(focal_lengths)
        # create pytorch mesh
        if color is None or color == 0:
            mesh_color = self.mesh_color_0
        else:
            mesh_color = self.mesh_color_1           
        tex = torch.ones_like(vertices) * mesh_color    # (1, V, 3)
        textures = Textures(verts_rgb=tex)
        mesh = Meshes(verts=vertices, faces=faces, textures=textures)
        # silhouette renderer 
        renderer_silh = MeshRenderer(
            rasterizer=MeshRasterizer(
                cameras=cameras, 
                raster_settings=self.raster_settings_soft),
            shader=SoftSilhouetteShader(blend_params=self.blend_params))
        # project silhouette
        silh_images = renderer_silh(mesh)[..., -1].unsqueeze(1)
        # project points
        if points is None:
            proj_points = None
        else:
            proj_points = self.project_points(points=points, cameras=cameras)
        if color is not None:
            # color renderer for visualization (no gradients)
            visualization = self._get_visualization_from_mesh(mesh, cameras)    
            return silh_images, proj_points, visualization
        else:
            return silh_images, proj_points