import numpy as np
from bvh import *

class Geometry(object):
    """
    Geometry object.
    """

    def __init__(self):
        self.materials = []
        self.mesh = np.empty(0)
        self.inside = np.empty(0)
        self.outside = np.empty(0)

    def add_mesh(self, mesh, inside, outside):
        """
        Add a triangle mesh.

        Args:
            - mesh: array of shape (n,3,3)
                Array of vertices for a closed mesh with n triangles.

            - inside: Material
                Material inside the mesh.

            - outside: Material
                Material outside the mesh.
         """

        if len(mesh.shape) != 3 or mesh.shape[1] != 3 or mesh.shape[2] != 3:
            raise Exception('shape mismatch')

        if inside not in self.materials:
            self.materials.append(inside)
        if outside not in self.materials:
            self.materials.append(outside)

        self.inside.resize(self.inside.size+mesh.shape[0])
        self.inside[-mesh.shape[0]:] = np.tile(self.materials.index(inside), mesh.shape[0])

        self.outside.resize(self.outside.size+mesh.shape[0])
        self.outside[-mesh.shape[0]:] = np.tile(self.materials.index(outside), mesh.shape[0])

        self.mesh.resize((self.mesh.shape[0]+mesh.shape[0],3,3))
        self.mesh[-mesh.shape[0]:] = mesh

    def build(self, bits=16):
        order, zvalues = morton_order(self.mesh, bits)


        zvalues = zvalues[order]
        self.mesh = self.mesh[order]
        self.inside = self.inside[order]
        self.outside = self.outside[order]

        leafs = []

        for z in sorted(set(zvalues)):
            mask = (zvalues == z)
            leafs.append(Leaf(self.mesh[mask], 3*np.where(mask)[0][0], z))

        layers = []
        layers.append(leafs)

        while True:
            zvalues = np.array([node.zvalue for node in layers[-1]])

            bit_shifted_zvalues = zvalues >> 1

            nodes = []
            for z in sorted(set(bit_shifted_zvalues)):
                mask = (bit_shifted_zvalues == z)
                nodes.append(Node(list(compress(layers[-1], mask)), z))

            layers.append(nodes)

            if len(nodes) == 1:
                break

        layers.reverse()

        nodes = []
        for layer in layers:
            for node in layer:
                nodes.append(node)


        self.lower_bound = np.empty((len(nodes),3))
        self.upper_bound = np.empty((len(nodes),3))
        self.child_map = np.empty(len(nodes))
        self.child_len = np.empty(len(nodes))
        self.first_leaf = -1

        for i, node in enumerate(nodes):
            self.lower_bound[i] = node.lower_bound
            self.upper_bound[i] = node.upper_bound
            
            self.child_len[i] = node.size

            if isinstance(node, Node):
                for j, child in enumerate(nodes):
                    if child is node.children[0]:
                        self.child_map[i] = j
                        break

            if isinstance(node, Leaf):
                self.child_map[i] = node.mesh_idx

                if self.first_leaf == -1:
                    self.first_leaf = i