import numpy as np


from deodr import Camera, Scene3D
from deodr.triangulated_mesh import TriMesh


class RenderCore:
    def __init__(self, max_depth=0, depth_scale=1):
        self.mesh = None
        self.depthScale = depth_scale
        self.scene = Scene3D(0)
        self.set_max_depth(max_depth)
		
    def set_mesh(self, vertices, faces):
        self.mesh = TriMesh(
            faces, vertices
        )  # we do a copy to avoid negative stride not support by Tensorflow
		
        object_center = vertices.mean(axis=0)
        object_radius = np.max(np.std(vertices, axis=0))
        self.camera_center = object_center + np.array([-0.5, 0, 5]) * object_radius
		
        self.scene.set_mesh(self.mesh)

    def set_max_depth(self, max_depth):
        self.scene.max_depth = max_depth
        self.scene.set_background_color([max_depth])

    def set_depth_scale(self, depth_scale):
        self.depthScale = depth_scale

    def set_image(self, width, height, intrinsic, extrinsic):
        self.width = width
        self.height = height
		
        # focal = None
        # if focal is None:
            # focal = 2 * self.width

        # rot = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
        # trans = -rot.T.dot(self.camera_center)
        # intrinsic = np.array(
            # [[focal, 0, self.width / 2], [0, focal, self.height / 2], [0, 0, 1]]
        # )
        # extrinsic = np.column_stack((rot, trans))
		
        self.camera = Camera(
            extrinsic=extrinsic,
            intrinsic=intrinsic,
            width=self.width,
            height=self.height,
            distortion=None,
        )
        
    def render(self):
        depth_scale = 1 * self.depthScale
        depth = self.scene.render_depth(
            self.camera,
            depth_scale=depth_scale,
        )
        #depth = tf.clip_by_value(depth, 0, self.scene.max_depth)
        return depth


""" CAMERA """
def intrinsic():
	RES_X = 640
	RES_Y = 480
	f_mm             = 50 # blender default
	sensor_w_mm      = 36 # blender default
	sensor_h_mm = sensor_w_mm * RES_Y / RES_X

	fx_px = f_mm * RES_X / sensor_w_mm
	fy_px = f_mm * RES_Y / sensor_h_mm

	u = RES_X / 2
	v = RES_Y / 2

	return np.array([[fx_px, 0,     u],
					 [0,     fy_px, v],
					 [0,     0,     1]], np.float32)
				

def extrinsic(camLoc):
	R_w2bc = np.array([[0, 1, 0],
					   [0, 0, 1],
					   [1, 0, 0]], np.float32)
	
	T_w2bc = -1 * R_w2bc.dot(camLoc)
	
	R_bc2cv = np.array([[1,  0,  0],
						[0, -1,  0],
						[0,  0, -1]], np.float32)
	
	R_w2cv = R_bc2cv.dot(R_w2bc)
	T_w2cv = R_bc2cv.dot(T_w2bc)

	return np.concatenate((R_w2cv, T_w2cv[:,None]), axis=1)
	
def proj(camLoc):
	return intrinsic().dot(extrinsic(camLoc))


class Render:
    def __init__(self, max_depth=0, depth_scale=1, size=(640, 480)):
        self.render_core = RenderCore(max_depth=max_depth)
        self.size = size
        self.max_depth = max_depth

    def render(self, mesh, faces, cam):
        self.render_core.set_mesh(mesh, faces)
        self.render_core.set_image(self.size[0], self.size[1], intrinsic(), extrinsic(cam))
        depth = self.render_core.render()

        thresh = self.max_depth - 1
        mask = depth<thresh
        # dmin = depth.min()
        dmean = depth[mask].mean()
        # dmax = depth[mask].max()
        depth[mask] = depth[mask] - dmean
        depth[~mask] = 0.0
        mask = mask.astype(np.float32)

        return depth


from torch import nn
from pytorch3d.renderer import (
    look_at_view_transform,
    FoVPerspectiveCameras, 
    PerspectiveCameras,
    PointLights, 
    DirectionalLights, 
    Materials, 
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    SoftPhongShader,
    TexturesUV,
    TexturesVertex
)
from pytorch3d.structures import Meshes

class DepthRenderer(nn.Module):
    def __init__(self) -> None:
         super().__init__()



    def forward(self, verts, faces, cam):
        R, T = look_at_view_transform(2.7, 0, 0) 
        cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T)
        ext = extrinsic(cam)
        cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=50)

        raster_settings = RasterizationSettings(
                            image_size=512, 
                            blur_radius=0.0, 
                            faces_per_pixel=1, 
                            bin_size=0
                        )

        rasterizer = MeshRasterizer(
                    cameras=cameras, 
                    raster_settings=raster_settings
                )
        depth = rasterizer(Meshes(verts[None,:,:], faces[None,:,:])).zbuf[0,:,:,0]

        mask = depth != -1
        # dmin = depth.min()
        dmean = depth[mask].mean()
        # dmax = depth[mask].max()
        depth[mask] = depth[mask] - dmean
        depth[~mask] = 0.0
        # mask = mask.astype(np.float32)

        return depth