import numpy as np
from itertools import chain

def compress(data, selectors):
    """
    Make an iterator that filters elements from `data` returning only those
    that have a corresponding element in `selectors` that evaluates to True.
    Stops when either the `data` or `selectors` iterables has been
    exhausted.

    From Python 2.7 itertools.
    """
    return (d for d, s in zip(data, selectors) if s)

class Leaf(object):
    """
    Leaf object represents the last layer in the bounding volume hierarchy
    which points to a subset of the triangle mesh.

    Args:
        - mesh: array,
            A subset of the triangle mesh.
        - mesh_idx: int,
            The index of the first triangle in the global mesh.
        - zvlaue: int, *optional*
            The zvalue associated with this leaf.

    .. note::
        The `mesh` passed in the constructor is not actually stored in the
        leaf. It is simply used to construct the leaf's bounding box.
    """
    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):
    """
    Node object represents a node in the bounding volume hierarchy which
    contains a list of child nodes.

    Args:
        - children: list,
            List of child nodes/leafs.
        - zvalue: int, *optional*
            The zvalue associated with this node.
    """
    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(arr):
    """
    Interleave the bits of quantized three-dimensional points in space.

    Example
        >>> interleave(np.identity(3, dtpe=np.int))
        array([4, 2, 1], dtype=uint64)
    """
    if len(arr.shape) != 2 or arr.shape[1] != 3:
        raise Exception('shape mismatch')

    z = np.zeros(arr.shape[0], dtype=np.uint64)
    for i in range(arr.dtype.itemsize*8):
        z |= (arr[:,2] & 1 << i) << (2*i) | \
             (arr[:,1] & 1 << i) << (2*i+1) | \
             (arr[:,0] & 1 << i) << (2*i+2)
    return z

def morton_order(mesh, bits):
    """
    Return a list of zvalues for triangles in `mesh` by interleaving the
    bits of the quantized center coordinates of each triangle. Each coordinate
    axis is quantized into 2**bits bins.
    """
    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))

    mean_positions = quantize(np.mean(mesh, axis=1))

    return interleave(mean_positions)

class Solid(object):
    """
    Object which stores a closed triangle mesh associated with a physically
    distinct object.

    Args:
        - mesh, array
            A closed triangle mesh.
        - material1, Material
            The material inside within the mesh.
        - material2, Material
            The material outside the mesh.
        - surface1, Surface,
            The surface on the inside of the mesh.
        - surface2, Surface,
            The surface on the outside of the mesh.

    .. warning::
        It is possible to define logically inconsistent geometries unless
        you are careful. For example, solid A may define its inside material
        as water but contain solid B which defines its outside material as air.
        In this case, a photon traveling out of solid B will reflect/refract
        assuming it's going into air, but upon reaching solid A's boundary will
        calculate attenuation factors assuming it just traveled through water.
    """
    def __init__(self, mesh, material1, material2, \
                     surface1=None, surface2=None):
        if len(mesh.shape) != 3 or mesh.shape[1] != 3 or mesh.shape[2] != 3:
            raise Exception('shape mismatch; mesh must be a triangle mesh')

        self.mesh = mesh
        self.material1 = material1
        self.material2 = material2
        self.surface1 = surface1
        self.surface2 = surface2

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

class Geometry(object):
    """Object which stores the global mesh for a geometry."""

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

    def add_solid(self, solid):
        """Add a solid to the geometry."""
        self.solids.append(solid)

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

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

        if solid.surface1 not in self.surfaces:
            self.surfaces.append(solid.surface1)

        if solid.surface2 not in self.surfaces:
            self.surfaces.append(solid.surface1)

    def build(self, bits=3):
        """Build the bounding volume hierarchy of the geometry."""
        mesh = np.concatenate([solid.mesh for solid in self.solids])

        # lookup solid/material/surface index from triangle index
        solid_index = \
            np.concatenate([np.tile(self.solids.index(solid), \
                                        len(solid)) for solid in self.solids])
        material1_index = \
            np.concatenate([np.tile(self.materials.index(solid.material1), \
                                        len(solid)) for solid in self.solids])
        material2_index = \
            np.concatenate([np.tile(self.materials.index(solid.material2), \
                                        len(solid)) for solid in self.solids])
        surface1_index = \
            np.concatenate([np.tile(self.surfaces.index(solid.surface1), \
                                        len(solid)) for solid in self.solids])
        surface2_index = \
            np.concatenate([np.tile(self.surfaces.index(solid.surface2), \
                                        len(solid)) for solid in self.solids])

        # array of zvalue for each triangle in mesh
        zvalues = morton_order(mesh, bits)

        # mesh indices in order of increasing zvalue
        order = np.array(zip(*sorted(zip(zvalues, range(zvalues.size))))[-1])

        # reorder all arrays ordered by triangle index
        self.zvalues = zvalues[order]
        self.mesh = mesh[order]
        self.solid_index = solid_index[order]
        self.material1_index = material1_index[order]
        self.material2_index = material2_index[order]
        self.surface1_index = surface1_index[order]
        self.surface2_index = surface2_index[order]

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

        layers = []
        layers.append(leafs)

        while True:
            bit_shifted_zvalues = \
                np.array([node.zvalue for node in layers[-1]]) >> 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 = None

        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 is None:
                    self.first_leaf = i