from ccpi.segmentation.simpleflex import ContourTree
import numpy

class SimpleflexSegmentor():
    

    PRIORITY_HEIGHT = 1
    PRIORITY_VOLUME = 2
    PRIORITY_HYPERVOLUME = 3
    PRIORITY_APPROX_HYPERVOLUME = 4
    
    ISO = 128
    LOWER = 127
    HIGHER = 129
    
    
    def __init__(self):
        self.dimensions = None
        self.isoValue = -1
        self.ct = ContourTree()
        
        self.scale = 1
        self.shift = 0
        
        self.origin = (0,0,0)
        self.spacing = (1.,1.,1.)
        
        self.isoVolume = 0
        self.highVolume = 0
        self.lowVolume = 0
        
        self.collapsePriority = 1
        

    def setImageDataDimensions(self, dimensions):
        isOk = False
        if type(dimensions) == tuple:
            if len(dimensions) == 3:
                isOk = True
        elif type(dimensions) == list:
            if len(dimensions) == 3:
                isOk = True
        elif type(dimensions) == numpy.ndarray:
            if len(dimensions) == 3:
                isOk = True
        if isOk:
            self.dimensions = numpy.asarray(dimensions)
    
    
    def setImageDataAsNumpyArray(self, img_data):
        '''Pass image data to the algorithm
        
        Performs a check of type as the algorithm accepts only unsigned char or
        unsigned short.'''
        
        dtype = img_data.dtype
        
        iMin = img_data.min()
        iMax = img_data.max()
        UNSIGNED_SHORT_MAX = 2 ** 16
        #UNSIGNED_CHAR_MAX = 2 ** 8
        sampleSize = 1
        # check the type
        if dtype == numpy.dtype('uint8'):
            # unsigned char => pass it as is
            self.scale = 1
            self.shift = 0
            sampleSize = 1
            #img_data_int = img_data.astype(numpy.dtype('uint8'))
        elif dtype == numpy.dtype('int8'):
            # signed char => should shift
            self.shift = -iMin
            img_data = img_data + self.shift
            sampleSize = 1
            #img_data_int = img_data.astype(numpy.dtype('uint8'))
        elif dtype == numpy.dtype('uint16'):
            # unsigned short => pass it as is
            self.scale = 1
            self.shift = 0
            sampleSize = 2
            #img_data_int = img_data.astype(numpy.dtype('uint16'))
        elif dtype == numpy.dtype('int16'):
            # signed short => shift
            self.shift = -iMin
            img_data = img_data + self.shift
            sampleSize = 2
            #img_data_int = img_data.astype(numpy.dtype('uint16'))
        else:
            # if the type is not any of the above
            # the image must be rescaled to unsigned short
            sampleSize = 2
            if (iMax - iMin != 0):
                self.scale = UNSIGNED_SHORT_MAX / (iMax - iMin)
            else:
                self.scale = 1
            self.shift = iMin
            img_data = img_data * self.scale
            img_data = img_data + self.shift
            #img_data_int = img_data.astype(numpy.dtype('uint16'))
                
        # the image is now correctly scaled to be passed to the algorithm
        
        # this conversion is ALWAYS LOSSLESS as we promote from SHORT (int16)
        # or CHAR (int8) to int32
        img_data_int = img_data.astype(numpy.int32)
        
        # pass it to the algorithm. 
        self.ct.CreateImageData(self.dimensions.astype(numpy.int32), sampleSize)
        self.ct.SetImageDataAsNumpyArray(img_data_int)
        
    def setIsoValuePercent(self, value):
        self.ct.SetIsoValuePercent(value)
        
    def setLocalIsoValuePercent(self, value):
        self.ct.SetLocalIsoValuePercent(value/100.)
    
    def setIsoValue(self, value):
        self.ct.SetIsoValue(value)

    def calculateContourTree(self):
        self.ct.CalculateContourTree()

    def constructIsoSurfaces(self):
        self.ct.ConstructIsoSurface()
    
    def constructLocalIsoSurfaces(self):
        self.ct.ConstructLocalIsoSurface()

    def getTrianglePoints(self):
        '''Returns the list of points belonging to the triangles
        
        The list contains a stream of x,y,z (image) coordinates of the points. 
        One point in 3D is defined by every 3 coordinates. 
        '''
        return self.ct.GetTrianglePoints()
    
    def getSurfacesList(self):
        '''Returns an index of the points belonging to a particular iso-surface'''
        return self.ct.GetSurfaces()

    def getSurfaces(self, surfno = None):
        '''Returns a numpy array with the points of each isosurface 
        
        Each points identified by a vector of 3 coordinates. This method returns
        an augmented vector, so that it can be used with augmented matrices to 
        perform affine transformations as scaling and translation at once.
        point = array([coordX, coordY, coordZ, 1])
        Each surface is identified by a number of points
        surface = array([point_0, point_1, point_2,... ])
        All the isosurfaces are put in an array, and that is what this method
        returns:
        listOfIsosurfaces = array([surface_0, surface_1, surface_2,...])
        '''
        
        if surfno == None:
            sorted_isosurface = self.ct.GetSurfaces()
        else:
            sorted_isosurface = [ self.ct.GetSurface(surfno) ]
            
        coord_list = self.ct.GetTrianglePoints()
        
        all_coords = []
        
        for t in sorted_isosurface:
            begin = t[0];
            end = t[1];
            i = begin
            
            surf_coord = []
            while i < end :
                x = numpy.asarray((coord_list[i + 0],
                                   coord_list[i + 1],
                                   coord_list[i + 2],
                                   1))
                surf_coord.append(x)
                i += 3
            all_coords.append(surf_coord)
        return numpy.asarray(all_coords)
            
        
    def setInputData(self, numpy3DArray):
        '''Set the 3D image data as numpy array

        The input array should have a shape (dimX,dimY,dimZ). The dimensions
        are calculated from it.'''
        shape = numpy.shape(numpy3DArray)
        if len(shape) == 3:
            self.setImageDataDimensions(shape)
            print(self.dimensions)
            # set the data shape as single array
            size = shape[0] * shape[1] * shape[2]
            self.setImageDataAsNumpyArray(numpy.reshape(numpy3DArray , (size,)))
        else:
            raise Exception("Input Data Dimensions Error: " + str(numpy.shape(numpy3DArray)))
            
    def setOrigin(self, origin):
        self.origin = list(origin)
    
    def setSpacing(self, spacing):
        self.spacing = list(spacing)
    
    def getSurfaceAsMask(self, which, labelIso=1, labelHigh=0, labelLow=0):            
        self.mask , counterIso, counterHigh, counterLow = self.ct.GetSurfaceAsMask(which, labelIso, labelHigh, labelLow)
        voxVolume = self.spacing[0] * self.spacing[1] * self.spacing[2]
        self.isoVolume = voxVolume * counterIso
        self.highVolume = voxVolume * counterHigh
        self.lowVolume = voxVolume * counterLow
        return self.mask
        
    def getSegmentedVolume(self, which=ISO):
        if which == 1:
            return self.isoVolume
        elif which == 2:
            return self.highVolume
        elif which == 3:
            return self.lowVolume
        
    def setCollapsePriority(self, priority=PRIORITY_HEIGHT):
        self.ct.SetCollapsePriority(priority)
        
    def resetCollapsePriority(self, priority=PRIORITY_HEIGHT):
        self.ct.ResetCollapsePriority(priority)
        
    def getCollapsePriority(self):
        return self.ct.GetCollapsePriority()
		
    def updateTreeFromLogPriorityBound(self, bound, isGlobal=True):
        if bound >=0 and bound <=1:
            self.ct.ResetCollapsePriority(self.collapsePriority)
            self.ct.SetLogPriorityBound(bound)
            self.ct.UpdateTreeFromLogPriorityBound()
            if isGlobal:
                self.constructIsoSurfaces()
            else:
                self.constructLocalIsoSurfaces()
    
    def updateTreeFromLogTreeSize(self, size, isGlobal=True):
        if size >=0 and size <=1:
            self.ct.ResetCollapsePriority(self.collapsePriority)
            self.ct.SetLogTreeSize(size)
            self.ct.UpdateTreeFromLogTreeSize()
            if isGlobal:
                self.constructIsoSurfaces()
            else:
                self.constructLocalIsoSurfaces()
            
    def getContourTree(self):
        return self.ct.GetContourTree()