import numpy as np

class Leaf(object):
    def __init__(self, meshasdf, mesh_idx):
        self.xmin, self.ymin, self.zmin = meshasdf[0]
        self.xmax, self.ymax, self.zmax = meshasdf[0]

        for x, y, z in meshasdf[1:]:
            if x < self.xmin:
                self.xmin = x
            if x > self.xmax:
                self.xmax = x
            if y < self.ymin:
                self.ymin = y
            if y > self.ymax:
                self.ymax = y
            if z < self.zmin:
                self.zmin = z
            if z > self.zmax:
                self.zmax = z

        self.mesh_idx = mesh_idx
        self.size = meshasdf.size//3

class Node(object):
    def __init__(self, children):
        self.xmin = children[0].xmin
        self.xmax = children[0].xmax
        self.ymin = children[0].ymin
        self.ymax = children[0].ymax
        self.zmin = children[0].zmin
        self.zmax = children[0].zmax

        for child in children[1:]:
            if child.xmin < self.xmin:
                self.xmin = child.xmin
            if child.xmax > self.xmax:
                self.xmax = child.xmax
            if child.ymin < self.ymin:
                self.ymin = child.ymin
            if child.ymax > self.ymax:
                self.ymax = child.ymax
            if child.zmin < self.zmin:
                self.zmin = child.zmin
            if child.zmax > self.zmax:
                self.zmax = child.zmax

        self.children = children
        self.size = len(children)

class Graph(object):
    def __init__(self, mesh, triangles_per_leaf=2, children_per_node=4):
        leafs = []
        for i in range(0, mesh.size, triangles_per_leaf*3):
            leafs.append(Leaf(mesh[i:i+triangles_per_leaf*3], i))

        layers = []
        layers.append(leafs)

        while len(layers[-1]) > 1:
            nodes = []
            for i in range(0, len(layers[-1]), children_per_node):
                nodes.append(Node(layers[-1][i:i+children_per_node]))

            layers.append(nodes)

        # leaves last
        layers.reverse()

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

        lower = []
        upper = []
        start = []
        count = [] # number of child nodes or triangles
        first_leaf = -1

        for i, node in enumerate(nodes):
            lower.append([node.xmin, node.ymin, node.zmin])
            upper.append([node.xmax, node.ymax, node.zmax])
            count.append(node.size)

            if isinstance(node, Node):
                for j, child in enumerate(nodes):
                    if child is node.children[0]:
                        start.append(j)
                        break

            if isinstance(node, Leaf):
                start.append(node.mesh_idx)

                if first_leaf == -1:
                    first_leaf = i

        self.lower = np.array(lower)
        self.upper = np.array(upper)
        self.start = np.array(start)
        self.count = np.array(count)
        self.first_leaf = first_leaf