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.mesh_idx = mesh_idx
        self.size = mesh.shape[0]
        self.zvalue = zvalue
        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.center = (self.lower_bound + self.upper_bound)/2.0

class Node(object):
    def __init__(self, children, zvalue=None):
        self.size = len(children)
        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.children = children

        self.center = (self.lower_bound + self.upper_bound)/2.0

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=8):
    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])])
    
    def quantize(x):
        return np.uint64((x-lower_bound)*bits/(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)))

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

    return order, zvalues