import numpy as np
from itertools import chain

def compress(data, selectors):
    return (d for d, s in zip(data, selectors) if s)

class Leaf(object):
    def __init__(self, mesh, mesh_idx, zvalue=None):
        self.zvalue = zvalue
        self.mesh_idx = mesh_idx

        self.lower_bound = np.array([np.min(mesh[:,:,0]), np.min(mesh[:,:,1]), np.min(mesh[:,:,2])])
        self.upper_bound = np.array([np.max(mesh[:,:,0]), np.max(mesh[:,:,1]), np.max(mesh[:,:,2])])

        self.size = mesh.shape[0]

class Node(object):
    def __init__(self, children, zvalue=None):
        self.zvalue = zvalue

        lower_pts = np.array([child.lower_bound for child in children])
        upper_pts = np.array([child.upper_bound for child in children])

        self.lower_bound = np.array([np.min(lower_pts[:,0]), np.min(lower_pts[:,1]), np.min(lower_pts[:,2])])
        self.upper_bound = np.array([np.max(upper_pts[:,0]), np.max(upper_pts[:,1]), np.max(upper_pts[:,2])])

        self.size = len(children)

        self.children = children

def interleave(*args):
    return int("".join(chain(*zip(*[bin(x)[2:].zfill(x.nbytes*8) for x in args]))), 2)

def morton_order(mesh, bits=3):
    lower_bound = np.array([np.min(mesh[:,:,0]), np.min(mesh[:,:,1]), np.min(mesh[:,:,2])])
    upper_bound = np.array([np.max(mesh[:,:,0]), np.max(mesh[:,:,1]), np.max(mesh[:,:,2])])

    max_value = 2**bits - 1
    
    def quantize(x):
        return np.uint64((x-lower_bound)*max_value/(upper_bound-lower_bound))

    zvalues = np.empty(mesh.shape[0], dtype=np.uint64)
    for i, triangle in enumerate(mesh):
        zvalues[i] = interleave(*quantize(np.mean(triangle, axis=0)))

    return zvalues

class Solid(object):
    def __init__(self, mesh, inside, outside):
        if len(mesh.shape) != 3 or mesh.shape[1] != 3 or mesh.shape[2] != 3:
            raise Exception('shape mismatch')

        self.mesh = mesh
        self.inside = inside
        self.outside = outside

    def __len__(self):
        return self.mesh.shape[0]

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

    def __init__(self):
        self.solids = []
        self.materials = []

    def add_solid(self, solid):
        self.solids.append(solid)

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

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

    def build(self, bits=3):
        self.mesh = np.concatenate([solid.mesh for solid in self.solids])
        self.solid_index = np.concatenate([np.tile(self.solids.index(solid), len(solid)) for solid in self.solids])
        self.inside_index = np.concatenate([np.tile(self.materials.index(solid.outside), len(solid)) for solid in self.solids])
        self.outside_index = np.concatenate([np.tile(self.materials.index(solid.outside), len(solid)) for solid in self.solids])

        zvalues = morton_order(self.mesh, bits)

        order = np.array(zip(*sorted(zip(zvalues, range(zvalues.size))))[-1])

        zvalues = zvalues[order]
        self.mesh = self.mesh[order]
        self.solid_index = self.solid_index[order]
        self.inside_index = self.inside_index[order]
        self.outside_index = self.outside_index[order]

        leafs = []
        for z in sorted(set(zvalues)):
            mask = (zvalues == z)
            leafs.append(Leaf(self.mesh[mask], 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