from data import Obj, ProjectionGenerator, Projection
import numpy as np
import math
import scipy.misc as m
from numba import jit
import torch

class TopDownView():
    """
    Take a room, pre-render top-down views
    Of floor, walls and individual objects
    That can be used to generate the multi-channel views used in our pipeline
    """
    def __init__(self, height_cap=4.05, length_cap=6.05, size=512):
        #Padding by 0.05m to avoid problems with boundary cases
        """
        Parameters
        ----------
        height_cap (int): the maximum height (in meters) of rooms allowed, which will be rendered with
            a value of 1 in the depth channel. To separate the floor from empty spaces,
            floors will have a height of 0.5m. See zpad below
        length_cap (int): the maximum length/width of rooms allowed.
        size (int): size of the rendered top-down image

        Returns
        -------
        visualization (Image): a visualization of the rendered room, which is
            simply the superimposition of all the rendered parts
        (floor, wall, nodes) (Triple[torch.Tensor, torch.Tensor, list[torch.Tensor]): 
            rendered invidiual parts of the room, as 2D torch tensors
            this is the part used by the pipeline
        """
        self.size = size
        self.pgen = ProjectionGenerator(room_size_cap=(length_cap, height_cap, length_cap), \
                                        zpad=0.5, img_size=size)

    def render(self, room):
        projection = self.pgen.get_projection(room)
        
        visualization = np.zeros((self.size,self.size))
        nodes = []

        for node in room.nodes:
            modelId = node.modelId #Camelcase due to original json

            t = np.asarray(node.transform).reshape(4,4)

            o = Obj(modelId)
            t = projection.to_2d(t)
            o.transform(t)
            
            t = projection.to_2d()
            bbox_min = np.dot(np.asarray([node.xmin, node.zmin, node.ymin, 1]), t)
            bbox_max = np.dot(np.asarray([node.xmax, node.zmax, node.ymax, 1]), t)
            xmin = math.floor(bbox_min[0])
            ymin = math.floor(bbox_min[2])
            xsize = math.ceil(bbox_max[0]) - xmin + 1
            ysize = math.ceil(bbox_max[2]) - ymin + 1

            description = {}
            description["modelId"] = modelId
            description["transform"] = node.transform
            description["bbox_min"] = bbox_min
            description["bbox_max"] = bbox_max
            
            #Since it is possible that the bounding box information of a room
            #Was calculated without some doors/windows
            #We need to handle these cases
            if ymin < 0: 
                ymin = 0
            if xmin < 0: 
                xmin = 0

            rendered = self.render_object(o, xmin, ymin, xsize, ysize, self.size)
            description["height_map"] = torch.from_numpy(rendered).float()
            
            tmp = np.zeros((self.size, self.size))
            tmp[xmin:xmin+rendered.shape[0],ymin:ymin+rendered.shape[1]] = rendered
            visualization += tmp

            nodes.append(description)
        
        #Render the floor
        o = Obj(room.modelId+"f", room.house_id, is_room=True)
        t = projection.to_2d()
        o.transform(t)
        floor = self.render_object(o, 0, 0, self.size, self.size, self.size)
        visualization += floor
        floor = torch.from_numpy(floor).float()
    
        #Render the walls
        o = Obj(room.modelId+"w", room.house_id, is_room=True)
        t = projection.to_2d()
        o.transform(t)
        wall = self.render_object(o, 0, 0, self.size, self.size, self.size)
        visualization += wall
        wall = torch.from_numpy(wall).float()
        
        return (visualization, (floor, wall, nodes))
    
    @staticmethod
    @jit(nopython=True)
    def render_object_helper(triangles, xmin, ymin, xsize, ysize, img_size):
        result = np.zeros((img_size, img_size), dtype=np.float32)
        N, _, _ = triangles.shape

        for triangle in range(N):
            x0,z0,y0 = triangles[triangle][0]
            x1,z1,y1 = triangles[triangle][1]
            x2,z2,y2 = triangles[triangle][2]
            a = -y1*x2 + y0*(-x1+x2) + x0*(y1-y2) + x1*y2
            if a != 0:
                for i in range(max(0,math.floor(min(x0,x1,x2))), \
                               min(img_size,math.ceil(max(x0,x1,x2)))):
                    for j in range(max(0,math.floor(min(y0,y1,y2))), \
                                   min(img_size,math.ceil(max(y0,y1,y2)))):
                        x = i+0.5
                        y = j+0.5
                        s = (y0*x2 - x0*y2 + (y2-y0)*x + (x0-x2)*y)/a
                        t = (x0*y1 - y0*x1 + (y0-y1)*x + (x1-x0)*y)/a
                        if s < 0 and t < 0:
                            s = -s
                            t = -t
                        if 0 < s < 1 and 0 < t < 1 and s + t <= 1:
                            height = z0 *(1-s-t) + z1*s + z2*t
                            result[i][j] = max(result[i][j], height)

        return result[xmin:xmin+xsize, ymin:ymin+ysize]
    
    @staticmethod
    def render_object(o, xmin, ymin, xsize, ysize, img_size):
        """
        Render a cropped top-down view of object
        
        Parameters
        ----------
        o (list[triple]): object to be rendered, represented as a triangle mesh
        xmin, ymin (int): min coordinates of the bounding box containing the object,
            with respect to the full image
        xsize, ysze (int); size of the bounding box containing the object
        img_size (int): size of the full image
        """
        triangles = np.asarray(list(o.get_triangles()), dtype=np.float32)
        return TopDownView.render_object_helper(triangles, xmin, ymin, xsize, ysize, img_size)
    
    @staticmethod
    @jit(nopython=True)
    def render_object_full_size_helper(triangles, size):
        result = np.zeros((size, size), dtype=np.float32)
        N, _, _ = triangles.shape

        for triangle in range(N):
            x0,z0,y0 = triangles[triangle][0]
            x1,z1,y1 = triangles[triangle][1]
            x2,z2,y2 = triangles[triangle][2]
            a = -y1*x2 + y0*(-x1+x2) + x0*(y1-y2) + x1*y2
            if a != 0:
                for i in range(max(0,math.floor(min(x0,x1,x2))), \
                               min(size,math.ceil(max(x0,x1,x2)))):
                    for j in range(max(0,math.floor(min(y0,y1,y2))), \
                                   min(size,math.ceil(max(y0,y1,y2)))):
                        x = i+0.5
                        y = j+0.5
                        s = (y0*x2 - x0*y2 + (y2-y0)*x + (x0-x2)*y)/a
                        t = (x0*y1 - y0*x1 + (y0-y1)*x + (x1-x0)*y)/a
                        if s < 0 and t < 0:
                            s = -s
                            t = -t
                        if 0 < s < 1 and 0 < t < 1 and s + t <= 1:
                            height = z0 *(1-s-t) + z1*s + z2*t
                            result[i][j] = max(result[i][j], height)
        
        return result

    @staticmethod
    def render_object_full_size(o, size):
        """
        Render a full-sized top-down view of the object, see render_object
        """
        triangles = np.asarray(list(o.get_triangles()), dtype=np.float32)
        return TopDownView.render_object_full_size_helper(triangles, size)