import numpy as np
from ...FiberVis_core import ROIsSegmentationQueryOctree, ROIsSegmentationPopulateAndDefragmentPool


class OctreePointBased:
    OT_NODE_CAPACITY = 256

    def __init__(self, points, fiberSizes):
        print('Building Octree...')
        self.vertex = points
        self.treeCount = 0
        self.root = 0
        self.pool = None
        self.vertexIdxPool = None

        x = self.vertex[0::3]
        y = self.vertex[1::3]
        z = self.vertex[2::3]

        xmin, xmax = x.min(), x.max()
        ymin, ymax = y.min(), y.max()
        zmin, zmax = z.min(), z.max()

        radius = np.array([(xmax-xmin)/2, (ymax-ymin)/2,
                          (zmax-zmin)/2], dtype=np.float32) * 1.005
        center = radius + np.array([xmin, ymin, zmin], dtype=np.float32)

        self.CPopulateAndDefragment(points.size, center, radius)

        self.createVertex2Fiber(fiberSizes)


    def CPopulateAndDefragment(self, vertexSize, center, radius):
        poolSize = int(vertexSize/self.OT_NODE_CAPACITY*4.4 + 1)

        dt = np.dtype([('children', np.uint32, (8,)),
                       ('numberOfElements', np.uint32, (1,)),
                       ('contentElements', np.uint32, (1,)),
                       ('radius', np.float32, (3,)),
                       ('center', np.float32, (3,))])

        self.pool = np.empty(poolSize, dtype=dt)
        self.vertexIdxPool = np.empty(self.vertex.size//3, dtype=np.uint32)

        n = ROIsSegmentationPopulateAndDefragmentPool(
            self.pool,
            poolSize,
            self.vertex,
            self.vertex.size,
            center,
            radius,
            self.vertexIdxPool)

        # n = cfuncs.ROIsSegmentationPopulateAndDefragmentPool(
        #       self.pool.ctypes.data,
        #       poolSize,
        #       self.vertex.ctypes.data,
        #       self.vertex.size,
        #       center.ctypes.data,
        #       radius.ctypes.data,
        #       self.vertexIdxPool.ctypes.data)

        self.pool = np.array(self.pool[:n], copy=True)

    def createVertex2Fiber(self, fiberSizes):
        self.vertexIdx2FiberIdx = np.empty(
            self.vertexIdxPool.size, dtype=np.int32)

        offset = 0
        for i in range(len(fiberSizes)):
            self.vertexIdx2FiberIdx[offset:offset+fiberSizes[i]] = i
            offset += fiberSizes[i]

    def queryCollision(self, queryPackage, result):
        ROIsSegmentationQueryOctree(len(queryPackage),
                                    queryPackage,
                                    self.pool,
                                    self.vertexIdxPool,
                                    self.vertexIdx2FiberIdx,
                                    self.vertex,
                                    result,
                                    result.shape[1])

        # cfuncs.ROIsSegmentationQueryOctree(   len(queryPackage),
        #                                                                       queryPackage.ctypes.data,
        #                                                                       self.pool.ctypes.data,
        #                                                                       self.vertexIdxPool.ctypes.data,
        #                                                                       self.vertexIdx2FiberIdx.ctypes.data,
        #                                                                       self.vertex.ctypes.data,
        #                                                                       result.ctypes.data,
        #                                                                       result.shape[1])
