class TestMCE(unittest.TestCase):
    def setUp(self):
        self.factor = Raster('../../examples/multifact.tif')
            #~ [1,1,3]
            #~ [3,2,1]
            #~ [0,3,1]

        self.state  = Raster('../../examples/sites.tif')
        self.state.resetMask(maskVals= [0])
            #~ [1,2,1],
            #~ [1,2,1],
            #~ [0,1,2]
        self.areaAnalyst = AreaAnalyst(self.state, second=None)

    def test_MCE(self):

        data = [
            [1.0,     4.0, 6.0, 7.0],
            [1.0/4,   1.0, 3.0, 4.0],
            [1.0/6, 1.0/3, 1.0, 2.0],
            [1.0/7, 1.0/4, 1.0/2, 1]
        ]
        # Multiband
        factor = Raster('../../examples/two_band.tif')

        mce = MCE([self.factor, factor, self.factor], data, 1, 2, self.areaAnalyst)
        w = mce.getWeights()
        answer = [0.61682294, 0.22382863, 0.09723423, 0.06211421]
        assert_almost_equal(w, answer)

        # One-band
        mce = MCE([self.factor, self.factor, self.factor, self.factor], data, 1, 2, self.areaAnalyst)
        w = mce.getWeights()
        answer = [0.61682294, 0.22382863, 0.09723423, 0.06211421]
        assert_almost_equal(w, answer)

        mask = [
            [False, False, False],
            [False, False, False],
            [False, False, True]
        ]
        p = mce.getPrediction(self.state).getBand(1)
        self.assertEquals(p.dtype, np.uint8)
        answer = [      # The locations where the big numbers are stored must be masked (see mask and self.state)
            [1,   3, 1],
            [1,   3, 1],
            [100, 1, 100]
        ]
        answer = np.ma.array(data = answer, mask = mask)
        assert_almost_equal(p, answer)

        c = mce.getConfidence().getBand(1)
        self.assertEquals(c.dtype, np.uint8)
        answer = [      # The locations where the big numbers are stored must be masked (see mask and self.state)
            [sum((w*100).astype(int))/3,   0,  sum((w*100).astype(int))],
            [sum((w*100).astype(int)),     0,  sum((w*100).astype(int))/3],
            [10000,  sum((w*100).astype(int)),  10000]
        ]
        answer = np.ma.array(data = answer, mask = mask)
        assert_almost_equal(c, answer)

    def sim(self):
        """
        Make 1 iteracion of simulation.
        """
        # TODO: eleminate AreaAnalyst.getChangeMap() from the process

        transition = self.crosstable.getCrosstable()

        prediction = self.getPrediction()
        state = self.getState()
        new_state = state.getBand(1).copy()  # New states (the result of simulation) will be stored there.
        analyst = AreaAnalyst(state, prediction)
        classes = analyst.classes
        changes = analyst.getChangeMap().getBand(1)

        # Make transition between classes according to
        # number of moved pixel in crosstable
        self.rangeChanged.emit(self.tr("Simulation process %p%"), len(classes) ** 2 - len(classes))
        for initClass in classes:
            for finalClass in classes:
                if initClass == finalClass:
                    continue

                # TODO: Calculate number of pixels to be moved via TransitoionMatrix and state raster
                n = transition.getTransition(
                    initClass, finalClass
                )  # Number of pixels to be moved (constant count now).
                # Find n appropriate places for transition initClass -> finalClass
                class_code = analyst.encode(initClass, finalClass)
                places = changes == class_code  # Array of places where transitions initClass -> finalClass are occured
                placesCount = np.sum(places)
                if placesCount < n:
                    self.logMessage.emit(
                        self.tr("There are more transitions in the transition matrix, then the model have found")
                    )
                    n = placesCount

                confidence = self.getConfidence().getBand(1)
                confidence = (
                    confidence * places
                )  # The higher is number in cell, the higer is probability of transition in the cell
                indices = []
                for i in range(n):
                    index = np.unravel_index(
                        confidence.argmax(), confidence.shape
                    )  # Select the cell with biggest probability
                    indices.append(index)
                    confidence[index] = -1  # Mark the cell to prevent second selection

                # Now "indices" contains indices of the appropriate places,
                # make transition initClass -> finalClass
                for index in indices:
                    new_state[index] = finalClass
                self.updateProgress.emit()

        result = Raster()
        result.create([new_state], state.getGeodata())
        self.state = result
        self.updatePrediction(result)
        self.processFinished.emit()

    def setUp(self):
        self.r1 = Raster('examples/multifact.tif')
        self.r2 = Raster('examples/sites.tif')
        self.r3 = Raster('examples/two_band.tif')

        # r1
        data1 = np.array(
            [
                [1,1,3],
                [3,2,1],
                [0,3,1]
            ])
        # r2
        data2 = np.array(
            [
                [1,2,1],
                [1,2,1],
                [0,1,2]
            ])
        mask = [
            [False, False, False],
            [False, False, False],
            [False, False, False]
        ]
        self.data1 = ma.array(data=data1, mask=mask)
        self.data2 = ma.array(data=data2, mask=mask)

class TestSimulator(unittest.TestCase):
    def setUp(self):

        # Raster1:
        # ~ [1, 1, 3,],
        # ~ [3, 2, 1,],
        # ~ [0, 3, 1,]
        self.raster1 = Raster("../../examples/multifact.tif")
        self.raster1.resetMask([0])

        self.X = np.array([[1, 2, 3], [3, 2, 1], [0, 1, 1]])
        self.X = np.ma.array(self.X, mask=(self.X == 0))
        self.raster2 = Raster()
        self.raster2.create([self.X], self.raster1.getGeodata())

        self.aa = AreaAnalyst(self.raster1, self.raster2)

        self.crosstab = CrossTableManager(self.raster1, self.raster2)

        # Simple model
        self.model = Model(state=self.raster1)

    def test_compute_table(self):

        # print self.crosstab.getCrosstable().getCrosstable()
        # CrossTab:
        #  [[ 3.  1.  0.]
        #   [ 0.  1.  0.]
        #   [ 1.  0.  2.]]
        prediction = self.model.getPrediction(self.raster1)
        # print prediction.getBand(1)
        # prediction = [[1.0 1.0 6.0]
        #  [6.0 5.0 1.0]
        #  [-- 6.0 1.0]]
        # confidence = self.model.getConfidence()
        # print confidence.getBand(1)
        # confidence =     [[1.0 0.5  0.33]
        #  [0.5 0.33 0.25]
        #  [--  0.25 0.2]]
        result = np.array([[2.0, 1.0, 3.0], [1.0, 2.0, 1.0], [0, 3.0, 1.0]])
        result = np.ma.array(result, mask=(result == 0))

        simulator = Simulator(
            state=self.raster1, factors=None, model=self.model, crosstable=self.crosstab
        )  # The model does't use factors
        simulator.setIterationCount(1)
        simulator.simN()
        state = simulator.getState().getBand(1)
        assert_array_equal(result, state)

        result = np.array([[2.0, 1.0, 1.0], [2.0, 2.0, 1.0], [0, 3.0, 1.0]])
        result = np.ma.array(result, mask=(result == 0))

        simulator = Simulator(self.raster1, None, self.model, self.crosstab)
        simulator.setIterationCount(2)
        simulator.simN()
        state = simulator.getState().getBand(1)
        assert_array_equal(result, state)

    def _predict(self, state, calcTransitions=False):
        """
        Predict the changes.
        """
        try:
            self.rangeChanged.emit(self.tr("Initialize model %p%"), 1)

            rows, cols = self.geodata["ySize"], self.geodata["xSize"]
            if not self.changeMap.geoDataMatch(state):
                raise WoeManagerError("Geometries of the state and changeMap rasters are different!")

            prediction = np.zeros((rows, cols), dtype=np.uint8)
            confidence = np.zeros((rows, cols), dtype=np.uint8)
            mask = np.zeros((rows, cols), dtype=np.byte)

            stateBand = state.getBand(1)

            self.updateProgress.emit()
            self.rangeChanged.emit(self.tr("Prediction %p%"), rows)

            for r in xrange(rows):
                for c in xrange(cols):
                    oldMax, currMax = -1000, -1000  # Small numbers
                    indexMax = -1  # Index of Max weight
                    initCat = stateBand[r, c]  # Init category (state before transition)
                    try:
                        codes = self.analyst.codes(initCat)  # Possible final states
                        for code in codes:
                            try:  # If not all possible transitions are presented in the changeMap
                                map = self.woe[code]  # Get WoE map of transition 'code'
                            except KeyError:
                                continue
                            w = map[r, c]  # The weight in the (r,c)-pixel
                            if w > currMax:
                                indexMax, oldMax, currMax = code, currMax, w
                        prediction[r, c] = indexMax
                        confidence[r, c] = int(100 * (sigmoid(currMax) - sigmoid(oldMax)))
                    except ValueError:
                        mask[r, c] = 1
                self.updateProgress.emit()

            predicted_band = np.ma.array(data=prediction, mask=mask, dtype=np.uint8)
            self.prediction = Raster()
            self.prediction.create([predicted_band], self.geodata)
            confidence_band = np.ma.array(data=confidence, mask=mask, dtype=np.uint8)
            self.confidence = Raster()
            self.confidence.create([confidence_band], self.geodata)
        except MemoryError:
            self.errorReport.emit(self.tr("The system out of memory during WOE prediction"))
            raise
        except:
            self.errorReport.emit(self.tr("An unknown error occurs during WoE prediction"))
            raise
        finally:
            self.processFinished.emit()

    def test_cat2vect(self):
        smp = Sampler(self.state, self.factors,  self.output, ns=0)
        assert_array_equal(smp.cat2vect(0), [1, 0])
        assert_array_equal(smp.cat2vect(1), [0, 1])
        assert_array_equal(smp.cat2vect(2), [0, 0])

        inputRast = Raster('../../examples/sites.tif')
        inputRast.resetMask([0])
        smp = Sampler(inputRast, self.factors,  self.output, ns=0)
        assert_array_equal(smp.cat2vect(1), [1])
        assert_array_equal(smp.cat2vect(2), [0])

    def train(self):
        """
        Train the model
        """
        self.transitionPotentials = {}
        try:
            iterCount = len(self.codes) * len(self.factors)
            self.rangeChanged.emit(self.tr("Training WoE... %p%"), iterCount)
            changeMap = self.changeMap.getBand(1)
            for code in self.codes:
                sites = binaryzation(changeMap, [code])
                # Reclass factors (continuous factor -> ordinal factor)
                wMap = np.ma.zeros(changeMap.shape)  # The map of summary weight of the all factors
                self.weights[code] = {}  # Dictionary for storing wheights of every raster's band
                for k in xrange(len(self.factors)):
                    fact = self.factors[k]
                    self.weights[code][k] = {}  # Weights of the factor
                    factorW = self.weights[code][k]
                    if self.bins:  # Get bins of the factor
                        bin = self.bins[k]
                        if (bin != None) and fact.getBandsCount() != len(bin):
                            raise WoeManagerError("Count of bins list for multiband factor is't equal to band count!")
                    else:
                        bin = None
                    for i in range(1, fact.getBandsCount() + 1):
                        band = fact.getBand(i)
                        if bin and bin[i - 1]:  #
                            band = reclass(band, bin[i - 1])
                        band, sites = masks_identity(band, sites, dtype=np.uint8)  # Combine masks of the rasters
                        woeRes = woe(
                            band, sites, self.unit_cell
                        )  # WoE for the 'code' (initState->finalState) transition and current 'factor'.
                        weights = woeRes["map"]
                        wMap = wMap + weights
                        factorW[i] = woeRes["weights"]
                    self.updateProgress.emit()

                # Reclassification finished => set WoE coefficients
                self.woe[code] = wMap  # WoE for all factors and the transition code.

                # Potentials are WoE map rescaled to 0--100 percents
                band = (sigmoid(wMap) * 100).astype(np.uint8)
                p = Raster()
                p.create([band], self.geodata)
                self.transitionPotentials[code] = p
                gc.collect()
        except MemoryError:
            self.errorReport.emit("The system out of memory during WoE trainig")
            raise
        except:
            self.errorReport.emit(self.tr("An unknown error occurs during WoE trainig"))
            raise
        finally:
            self.processFinished.emit()

 def test_save(self):
     try:
         filename = 'temp.tiff'
         self.r1.save(filename)
         r2 = Raster(filename)
         self.assertEqual(r2.get_dtype(), self.r1.get_dtype())
         self.assertEqual(r2.getBandsCount(), self.r1.getBandsCount())
         for i in range(r2.getBandsCount()):
             assert_array_equal(r2.getBand(i+1), self.r1.getBand(i+1))
     finally:
         os.remove(filename)

 def errorMap(self, answer):
     '''
     Create map of correct and incorrect prediction.
     This function compares the known answer and the result of predicting procedure,
     correct pixel is marked as 0.
     '''
     state = self.getState()
     b = state.getBand(1)
     a = answer.getBand(1)
     diff = (a-b).astype(np.int16)
     result = Raster()
     result.create([diff], state.getGeodata())
     return result

    def setUp(self):
        self.factor = Raster('../../examples/multifact.tif')
                #~ [1,1,3]
                #~ [3,2,1]
                #~ [0,3,1]

        self.sites  = Raster('../../examples/sites.tif')
                    #~ [1,2,1],
                    #~ [1,2,1],
                    #~ [0,1,2]
        self.sites.resetMask(maskVals= [0])

        self.mask = [
            [False, False, False,],
            [False, False, False,],
            [True,  False, False,]
        ]
        fact = [
            [1, 1, 3,],
            [3, 2, 1,],
            [0, 3, 1,]
        ]
        site = [
            [False, True,  False,],
            [False, True,  False,],
            [False, False, True,]
        ]
        self.factraster  = ma.array(data = fact, mask=self.mask, dtype=np.int)
        self.sitesraster = ma.array(data = site, mask=self.mask, dtype=np.bool)

    def setUp(self):
        self.factor = Raster('../../examples/multifact.tif')
            #~ [1,1,3]
            #~ [3,2,1]
            #~ [0,3,1]

        self.state  = Raster('../../examples/sites.tif')
        self.state.resetMask(maskVals= [0])

 def makeChangeMap(self):
     f, s = self.first, self.second
     rows, cols = self.geodata['ySize'], self.geodata['xSize']
     band = np.zeros([rows, cols])
     self.rangeChanged.emit(self.tr("Creating change map %p%"), rows)
     for i in xrange(rows):
         for j in xrange(cols):
             if not f.mask[i,j]:
                 r = f[i,j]
                 c = s[i,j]
                 band[i, j] = self.encode(r, c)
         self.updateProgress.emit()
     bands = [np.ma.array(data = band, mask = f.mask)]
     raster = Raster()
     raster.create(bands, self.geodata)
     self.processFinished.emit(raster)
     self.changeMap = raster

class TestLRManager (unittest.TestCase):
    def setUp(self):
        self.output  = Raster('../../examples/multifact.tif')
            #~ [1,1,3]
            #~ [3,2,1]
            #~ [0,3,1]
        self.output.resetMask([0])
        self.state   = self.output
        self.factors = [Raster('../../examples/sites.tif'), Raster('../../examples/sites.tif')]
            #~ [1,2,1],
            #~ [1,2,1],
            #~ [0,1,2]


        self.output1  = Raster('../../examples/data.tif')
        self.state1   = self.output1
        self.factors1 = [Raster('../../examples/fact16.tif')]

    def test_LR(self):
        data = [
            [3.0, 1.0, 3.0],
            [3.0, 1.0, 3.0],
            [0,   3.0, 1.0]
        ]
        result = np.ma.array(data = data, mask = (data==0))

        lr = LR(ns=0)   # 3-class problem
        lr.setTrainingData(self.state, self.factors, self.output)
        lr.train()
        predict = lr.getPrediction(self.state, self.factors)
        predict = predict.getBand(1)
        assert_array_equal(predict, result)

        lr = LR(ns=1) # Two-class problem (it's because of boundary effect)
        lr.setTrainingData(self.state1, self.factors1, self.output1)
        lr.train()
        predict = lr.getPrediction(self.state1, self.factors1)
        predict = predict.getBand(1)
        data = [
            [0.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 2.0, 0.0],
            [0.0, 2.0, 2.0, 0.0],
            [0.0, 0.0, 0.0, 0.0],
        ]
        result = np.ma.array(data = data, mask = (data==0))
        assert_array_equal(predict, result)

    def test_get_state(self):
        smp = Sampler(self.state, self.factors,  self.output, ns=0)
        assert_array_equal(smp.get_state(self.state, 1,1), [0, 0])
        assert_array_equal(smp.get_state(self.state, 0,0), [0, 1])

        smp = Sampler(self.state, self.factors,  self.output, ns=1)
        res = [ 0, 1, 0, 0, 0, 1,
                0, 1, 0, 0, 0, 1,
                1, 0, 0, 1, 0, 0,
            ]
        assert_array_equal(smp.get_state(self.state, 1,1), res)

        inputRast = Raster('../../examples/sites.tif')
        inputRast.resetMask([0])
        smp = Sampler(inputRast, self.factors,  self.output, ns=0)
        assert_array_equal(smp.get_state(self.state, 1,1), [0])
        assert_array_equal(smp.get_state(self.state, 0,0), [1])

class TestAreaAnalysisManager (unittest.TestCase):
    def setUp(self):
        self.r1 = Raster('../../examples/multifact.tif')
        # r1 -> r1 transition
        self.r1r1 = [
            [5,  5,  15,],
            [15, 10, 5, ],
            [0,  15, 5, ]
        ]
        
        self.r2  = Raster('../../examples/multifact.tif')
        self.r2.resetMask([0])
        self.r2r2 = [
            [0,   0, 8,],
            [8,   4, 0,],
            [100, 8, 0,]
        ]
        
        self.r3 = Raster('../../examples/multifact.tif')
        self.r3.resetMask([2])
        
    def test_AreaAnalyst(self):
        aa = AreaAnalyst(self.r1, self.r1)
        raster = aa.getChangeMap()
        band = raster.getBand(1)
        assert_array_equal(band, self.r1r1)
        
        # Masked raster
        aa = AreaAnalyst(self.r2, self.r2)
        raster = aa.getChangeMap()  
        band = raster.getBand(1)
        assert_array_equal(band, self.r2r2)
    
    def test_encode(self):
        aa = AreaAnalyst(self.r1, self.r1)
        self.assertEqual(aa.classes, [0,1,2,3])
        self.assertEqual(aa.encode(1,2), 6)
        for initClass in range(4):
            for finalClass in range(4):
                k = aa.encode(initClass, finalClass)
                self.assertEqual(aa.decode(k), (initClass, finalClass))
        self.assertEqual(aa.finalCodes(0), [0,1,2,3])
        self.assertEqual(aa.finalCodes(1), [4,5,6,7])

    def _predict(self, state):
        '''
        Predict the changes.
        '''
        geodata = self.changeMap.getGeodata()
        rows, cols = geodata['ySize'], geodata['xSize']
        if not self.changeMap.geoDataMatch(state):
            raise WoeManagerError('Geometries of the state and changeMap rasters are different!')

        prediction = np.zeros((rows,cols))
        confidence = np.zeros((rows,cols))
        mask = np.zeros((rows,cols))

        woe = self.getWoe()
        stateBand = state.getBand(1)

        for r in xrange(rows):
            for c in xrange(cols):
                oldMax, currMax = -1000, -1000  # Small numbers
                indexMax = -1                   # Index of Max weight
                initClass = stateBand[r,c]      # Init class (state before transition)
                try:
                    codes = self.analyst.codes(initClass)   # Possible final states
                    for code in codes:
                        try: # If not all possible transitions are presented in the changeMap
                            map = woe[code]     # Get WoE map of transition 'code'
                        except KeyError:
                            continue
                        w = map[r,c]        # The weight in the (r,c)-pixel
                        if w > currMax:
                            indexMax, oldMax, currMax = code, currMax, w
                    decode = self.analyst.decode(indexMax)    # Get init & final classes (initState, finalState)
                    prediction[r,c] = decode[1]               # final class
                    confidence[r,c] = sigmoid(currMax) - sigmoid(oldMax)
                except ValueError:
                    mask[r,c] = 1

        predicted_band = np.ma.array(data=prediction, mask=mask)
        self.prediction = Raster()
        self.prediction.create([predicted_band], geodata)
        confidence_band = np.ma.array(data=confidence, mask=mask)
        self.confidence = Raster()
        self.confidence.create([confidence_band], geodata)

 def setUp(self):
     self.r1 = Raster('../../examples/multifact.tif')
     # r1 -> r1 transition
     self.r1r1 = [
         [5,  5,  15,],
         [15, 10, 5, ],
         [0,  15, 5, ]
     ]
     
     self.r2  = Raster('../../examples/multifact.tif')
     self.r2.resetMask([0])
     self.r2r2 = [
         [0,   0, 8,],
         [8,   4, 0,],
         [100, 8, 0,]
     ]
     
     self.r3 = Raster('../../examples/multifact.tif')
     self.r3.resetMask([2])

    def _predict(self, state):
        '''
        Predict the changes.
        '''
        geodata = state.getGeodata()
        rows, cols = geodata['ySize'], geodata['xSize']

        # Get locations where self.initStateNum is occurs
        band = state.getBand(1)
        initStateMask = binaryzation(band, [self.initStateNum])
        mask = band.mask

        # Calculate summary map of factors weights
        # Confidence:
        #   confidence is summary map of factors, if current state = self.initState
        #   confidence is 0, if current state != self.initState
        # Prediction:
        #   predicted value is a constant = self.finalStateNum, if current state = self.initState
        #   predicted value is current state, if current state != self.initState
        confidence = np.zeros((rows,cols))
        weights = self.getWeights()
        weightNum = 0               # Number of processed weights
        for f in self.factors:
            if not f.geoDataMatch(state):
                raise MCEError('Geometries of the state and factor rasters are different!')
            f.normalize(mode = 'maxmin')
            for i in xrange(f.getBandsCount()):
                band = f.getBand(i+1)
                confidence = confidence + band*weights[weightNum]
                mask = np.ma.mask_or(mask, band.mask)
                weightNum = weightNum + 1
        confidence = confidence*initStateMask
        prediction = np.copy(state.getBand(1))
        prediction = np.logical_not(initStateMask) * prediction
        prediction = prediction + initStateMask*self.finalStateNum

        predicted_band = np.ma.array(data=prediction, mask=mask)
        self.prediction = Raster()
        self.prediction.create([predicted_band], geodata)
        confidence_band = np.ma.array(data=confidence, mask=mask)
        self.confidence = Raster()
        self.confidence.create([confidence_band], geodata)

    def setUp(self):

        # Raster1:
        # ~ [1, 1, 3,],
        # ~ [3, 2, 1,],
        # ~ [0, 3, 1,]
        self.raster1 = Raster("../../examples/multifact.tif")
        self.raster1.resetMask([0])

        self.X = np.array([[1, 2, 3], [3, 2, 1], [0, 1, 1]])
        self.X = np.ma.array(self.X, mask=(self.X == 0))
        self.raster2 = Raster()
        self.raster2.create([self.X], self.raster1.getGeodata())

        self.aa = AreaAnalyst(self.raster1, self.raster2)

        self.crosstab = CrossTableManager(self.raster1, self.raster2)

        # Simple model
        self.model = Model(state=self.raster1)

    def setUp(self):
        self.factor = Raster('../../examples/multifact.tif')
            #~ [1,1,3]
            #~ [3,2,1]
            #~ [0,3,1]

        self.state  = Raster('../../examples/sites.tif')
        self.state.resetMask(maskVals= [0])
            #~ [1,2,1],
            #~ [1,2,1],
            #~ [0,1,2]
        self.areaAnalyst = AreaAnalyst(self.state, second=None)