import math
import pickle 
import numpy as np
import matplotlib.pyplot as plt 
import matplotlib.ticker as ticker
import random as rand  
from numpy import random 
from sklearn import decomposition 
from deap import creator, base, tools, algorithms 
from sklearn.decomposition import PCA
from mpl_toolkits.mplot3d import Axes3D
from sklearn.cluster import KMeans  
import matplotlib.ticker as ticker 
import os.path 

from model_clb import model_clb, param_references



import time  
import multiprocessing

'''
Regions consist of cloud of points and principal component that govern the direction of exploration  
''' 
class Region: 
    def __init__(self, points, model, label, depth=1):    
        self.points = np.array(points)  
        self.model = model  
        self.pca = PCA(n_components=self.model.nParams)
        self.components = None
        self.prevComponents = None 
        self.cluster = False
        self.terminated = False  
        self.iter = 0      
        self.maxIter = 10            
        self.threshold = 0.001    
        self.label = label
        self.maxVarScale = 6      
        self.minVarScale = 3       
        self.varScaleDt = (self.maxVarScale - self.minVarScale)/(float(self.maxIter))                       
        self.varScale = self.maxVarScale         
        self.depth = depth     
        
    def updateVariance(self): 
        self.varScale = self.varScale - self.varScaleDt

    def updateIter(self):
        self.iter = self.iter + 1
        self.updateVariance()           
        
    def fitPCA(self): 
        self.prevComponents = self.components 
        self.pca.fit(self.points)
        self.components = self.pca.components_
    
    def transform(self, points):  
        return self.pca.transform(points)  
        
    def inverse_transform(self, points):
        return self.pca.inverse_transform(points)   
        
    def converged(self):
        if self.components is None or self.prevComponents is None: 
            return False        
        return np.linalg.norm(self.components - self.prevComponents) < self.threshold   
        
    def explored(self):    
        return self.terminated or self.iter > self.maxIter or self.converged()   
        
'''
The main class
'''

#must be global for parallelization     
creator.create("FitnessMax", base.Fitness, weights=(1.0,)) 
creator.create("Candidate", list, fitness=creator.FitnessMax)                                       
                                                           
                                                               

class Solver:
    def __init__(self, model, populationSize=10000, NGEN = 10, nsamples = 1e5, random_candidates = False):                                                                    
        self.model = model                   
        self.populationSize = populationSize          
        self.NGEN = NGEN              
        self.nsamples = int(nsamples)       
        self.indpb = 0.75           
        
        #GA operators
        #creator.create("FitnessMax", base.Fitness, weights=(1.0,)) 
        #creator.create("Candidate", list, fitness=creator.FitnessMax)       
        self.toolbox = base.Toolbox()    
        if random_candidates: 
            self.toolbox.register("candidate", self.generateRandomCandidate) 
        else:
            self.toolbox.register("candidate", self.generateCandidate) 
        self.toolbox.register("population", tools.initRepeat, list, self.toolbox.candidate)  
        self.toolbox.register("mate", tools.cxTwoPoint)
        self.toolbox.register("mutate", self.mutateCandidate, indpb=self.indpb, mult=0.5)      
        self.toolbox.register("select", tools.selTournament, tournsize=int(self.populationSize/10))             
    
    #estimate initial values with GA
    def findNominalValues(self):            

        tic = time.perf_counter()                         
        nominalVals = []   
        
        for evalMode in self.model.modes: 
            nominalValsMode = []
   
            #initialize new random population
            self.popu = self.toolbox.population(self.populationSize)     
            self.toolbox.register("evaluate", evalMode)   
                
            
            pool = multiprocessing.Pool()
            self.toolbox.register("map", pool.map)                             
                                                    
            
            
            for gen in range(self.NGEN):  
                print(gen)           
                #generate offspprings with crossover and mutations
                offspring = algorithms.varAnd(self.popu, self.toolbox, cxpb=0.5, mutpb=0.75)  
                #evaluate individuals
                fits = self.toolbox.map(self.toolbox.evaluate, offspring)     
                for fit, ind in zip(fits, offspring):      
                    if self.model.isViable(ind, fitness=fit) and ind not in nominalValsMode:      
                        nominalValsMode.append(ind)          
                    ind.fitness.values = fit      
                #roulete wheel selection
                self.popu = self.toolbox.select(offspring, k=len(self.popu))      
            
            pool.close()
            print("Number of viable points: " + str(len(nominalValsMode))) 
            nominalVals.extend(nominalValsMode)     
        toc = time.perf_counter()
        print("Elapsed time: " + str(toc - tic) + "s")                          
        return nominalVals        
        
    #creates an array of random candidates 
    def generateRandomCandidate(self): 
        candidate = []
        for ind in range(self.model.nParams): 
            candidate.append(random.uniform(self.model.parameter_values[self.model.params[ind]]["min"], self.model.parameter_values[self.model.params[ind]]["max"]))            
        return creator.Candidate(candidate) 
     
     
    def generateCandidate(self): 
        candidate = []
        for ind in range(self.model.nParams): 
            #candidate.append(random.uniform(self.model.parameter_values[self.model.params[ind]]["min"], self.model.parameter_values[self.model.params[ind]]["max"]))

            while True:
                r = random.normal(param_references[ind],np.min([
                                	np.abs(self.model.parameter_values[self.model.params[ind]]["min"] - param_references[ind]),
							        np.abs(self.model.parameter_values[self.model.params[ind]]["max"] - param_references[ind])
                                   ])/3.0)
                if r > 0:
                    break

            candidate.append(r)

            
        return creator.Candidate(candidate)
        
    def checkOutAllBounds(self, candidate):
        for idx, val in enumerate(candidate):
            if self.checkOutOfBounds(candidate, idx): 
                return True  
        return False      
                
    def checkOutOfBounds(self, candidate, idx): 
        #if out of bounds return True 
        if candidate[idx] < self.model.parameter_values[self.model.params[idx]]["min"] or candidate[idx] > self.model.parameter_values[self.model.params[idx]]["max"]: 
            return True
        return False            
    
    #returns a tuple of mutated candidate   
    def mutateCandidate(self, candidate, indpb, mult):  
        for idx, val in enumerate(candidate):   
            rnd = random.uniform(0, 1)
            if rnd >= indpb:
                rnd2 = random.uniform(1 - mult, 1 + mult)   
                candidate[idx] = val*rnd2   
                if candidate[idx] < self.model.parameter_values[self.model.params[idx]]["min"]: 
                    candidate[idx] = self.model.parameter_values[self.model.params[idx]]["min"]  
                if candidate[idx] > self.model.parameter_values[self.model.params[idx]]["max"]:  
                    candidate[idx] = self.model.parameter_values[self.model.params[idx]]["max"]                     
        return candidate,     
    
    def getViablePoints(self, points):
        pool = multiprocessing.Pool()

        viables = np.array(pool.map(self.model.isViable, points))  
        viable = np.array(points)[viables]
        """                
        viable = list() 
        i = 0
        for point in points:  
            i += 1
            if i % 1000 == 0:
                print(i)     
            
            #check if point is viable 
            if self.model.isViable(point): 
                viable.append(point)        
        """
        pool.close()                         
        return viable          
    
    # gap statistic method
    # returns the optimal number of clusters    
    def gapStatistic(self, region, number_ref = 10, max_clusters = 2, plot = False):        
        #sample size is equal to the number of samples in gaussian sampling  
        sample_size = self.nsamples    
        subjects = np.array(region.points)                 
        gaps = []
        deviations = []   
        references = [] 
        clusters_range = range(1, max_clusters + 1) 
        
        transformed = region.transform(subjects) 
        #get min and max parameter values in pca space 
        minP = np.min(transformed, axis=0)  
        maxP = np.max(transformed, axis=0)   
        
        for gap_clusters in clusters_range:
            print(gap_clusters) 
            reference_inertia = []  
            for index in range(number_ref): 

                #OBB ... orientated bounding box 
                #random sampling within the PCA bounding box            
                reference = minP + random.rand(sample_size, self.model.nParams)*(maxP - minP)
                reference = region.inverse_transform(reference) 
                
                kmeanModel = KMeans(gap_clusters) 
                kmeanModel.fit(reference) 
                reference_inertia.append(kmeanModel.inertia_)    
            
            kmeanModel = KMeans(gap_clusters)      
            kmeanModel.fit(subjects)     
            log_ref_inertia = np.log(reference_inertia)  
            #calculate gap
            gap = np.mean(log_ref_inertia) - np.log(kmeanModel.inertia_)  
            sk = math.sqrt(1 + 1.0/number_ref)*np.std(log_ref_inertia)  
            gaps.append(gap)    
            deviations.append(sk)                   
            
        # Plot the gaps         
        if plot:
            plt.clf() 
            ax = plt.gca() 
            ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True))     
            ax.xaxis.set_major_locator(ticker.MultipleLocator(2))   
            lines = plt.errorbar(clusters_range, gaps, ecolor='dodgerblue', yerr=deviations, fmt='-', color='dodgerblue') 
            plt.setp(lines[0], linewidth=1.5)  
            plt.ylabel('Gaps')
            plt.show()  
            
        #return optimal number of clusters
        for k in range(0, max_clusters - 1): 
            if gaps[k] >= gaps[k + 1] - deviations[k + 1]: 
                print("Optimal number of clusters: " + str(k + 1)) 
                return k + 1     
        print("Optimal number of clusters: " + str(max_clusters))       
        return max_clusters
    #returns the viable volume for 
    def getViableVolume(self, viableRegions, sample_size = int(1e4)): #1e4      
        volume = 0 

        for region in viableRegions:        
            regPoints = region.points
            region.fitPCA()                                                 
            transformed = region.transform(regPoints) 
                
            minP = np.min(transformed, axis=0)   
            maxP = np.max(transformed, axis=0)   
        
            dP = maxP - minP
            volB = np.prod(dP)          

            mcRef = minP + random.rand(sample_size, self.model.nParams)*dP  
            mcRef = region.inverse_transform(mcRef)         
            
            viaPoints = self.getViablePoints(mcRef) 
            count = np.ma.size(viaPoints, axis=0) 
            
            #volume for region  
            ratio = count/sample_size   
            volume = volume + ratio*volB            
    
        total = self.model.getTotalVolume()
        ratio = volume/total

        description = "Bounding box volume:" + str(volB) + "\n" 
        description += "Volume:" + str(volume) + "\n" 
        description += "Total volume:" + str(total) + "\n"    
        description +=  "Volume ratio:" + str(ratio)       
        
        return (volume, total, ratio), description    

    def setBoxColors(self, bp, nRegions, ax, colors = ["#0E74C8", "#15A357", "r", "k"]):
        colorLen = len(colors)

        for i in range(nRegions): 
            col = colors[i % colorLen]       
            plt.setp(bp['boxes'][i], color=col, linewidth=1.5)    
            plt.setp(bp['caps'][2*i], color=col, linewidth=1.5)  
            plt.setp(bp['caps'][2*i + 1], color=col, linewidth=1.5) 
            plt.setp(bp['whiskers'][2*i], color=col, linewidth=1.5)  
            plt.setp(bp['whiskers'][2*i + 1], color=col, linewidth=1.5)   
            plt.setp(bp['fliers'][i], color=col) 
            plt.setp(bp['medians'][i], color=col, linewidth=1.5)   
        
    def plotParameterVariances(self, viableSets, names=None, units=None):      
        #go through all parameters  
        params = self.model.params    
        figure = plt.figure()     
        nRows = math.ceil(len(params)/3)    
        for pcount, param in enumerate(params):    
            ax1 = plt.subplot(nRows, 3, pcount+1)  
            #if names == None:
            #   ax1.set_title(str(param) + str(pcount))    
            #else:
            #   ax1.set_title(names[pcount])  
            if units != None:
                plt.ylabel(names[pcount] + " " + units[pcount])  
            allRegions = []     
            #go through all regions 
            numSets = len(viableSets) 
            allNames = []
            allBoxes = []
            for count, reg in enumerate(viableSets): 
                points = np.array(reg.points)    
                data = points[:,pcount]   
                allRegions.append(data)   
                allNames.append("Region " + str(count + 1))                 
            bp = ax1.boxplot(allRegions, positions=list(range(1, numSets + 1)), widths = 0.4) 
            self.setBoxColors(bp, numSets, ax1)         
            allBoxes = bp['boxes'] 
            
        #draw legend 
        figure.legend(allBoxes, allNames, 'lower right')
        plt.show()     
        
    #Main method  
    def run(self, filename, maxDepth=0):    
        #filename is a file to which viable sets will be serialized    

        #estimate the inital viable set 
        viablePoints = self.findNominalValues()                                         
        
        if not viablePoints: 
            print("No viable points found!")  
            return 
        
        #dump viable points to file  
        pickle.dump(viablePoints, open(filename + "ViableSet_IterGA.p", "wb+"))   
        
        reg = Region(viablePoints, self.model, "0")   
        reg.fitPCA() 
        
        fpca = PCA(n_components=2)                  
        fpca.fit(reg.points)
                
        viableSets = list() 
        viableSets.append(reg)                   
        converged = False       
        iter = 0 
        
        while not converged: 
            converged = True                                                 
            iter += 1 
            print("Iteration: " + str(iter))    
            for set in viableSets:                  
                set.updateIter() 
                #if set not already explored  
                if not set.explored():
                    setSize = len(set.points) 
                    print("Label: " + set.label)   
                    print("Iter: " + str(set.iter))  
                    print("Variance scaling factor: " + str(set.varScale))                                  
                    converged = False                       
                          
                    #sample with 0 mean and scaled variance of prinicpal components       
                    candidateSet = random.multivariate_normal([0]*self.model.nParams, np.diag(set.pca.explained_variance_)*set.varScale, self.nsamples)             
                    candidateSet = set.inverse_transform(candidateSet)      
                                
                    #check if parameter values are not out of range         
                    inBounds = list() 
                    for cand in candidateSet:               
                        if not self.checkOutAllBounds(cand): 
                            inBounds.append(cand)  
                    inBounds = np.array(inBounds)           
                    candidateSet = inBounds   

                    Y = fpca.transform(candidateSet)                    
                    X = fpca.transform(set.points) 
                    
                    fig = plt.figure(iter)  
                    plt.clf()          
                    plt.scatter(Y[:, 0], Y[:, 1], c="red", alpha=0.1, edgecolor='k', rasterized=True)  
                    plt.scatter(X[:, 0], X[:, 1], c="cornflowerblue", alpha=0.8, edgecolor='k', rasterized=True)                
                    
                    plt.xlabel('PC 1')      
                    plt.ylabel('PC 2')    
                    plt.savefig(filename + "Set" + set.label + "Iter" + str(set.iter) + ".pdf")          
                    #identify viable points  
                    viablePoints = np.array(self.getViablePoints(candidateSet))  
                    
                    #if viable set is smaller than number of parameters do not accept it
                    print("Number of viable points: " + str(len(viablePoints)))                     
                    pickle.dump(candidateSet, open(filename + "_Region" + str(set.label) + "CandidateSet_Iter" + str(set.iter) +  ".p", "wb+"))     
                    pickle.dump(viablePoints, open(filename + "_Region" + str(set.label) + "ViableSet_Iter" + str(set.iter) + ".p", "wb+"))                         
                    set.points = viablePoints           
                    set.fitPCA()   
                        
                #if set not already terminated, terminate it and cluster   
                elif not set.terminated:     
                    set.terminated = True        
                    set.cluster = True                              
                        
            #clustering, check for new clusters                     
            newViableSets = list()   
            for set in viableSets: 
                if set.cluster and (maxDepth == 0 or set.depth < maxDepth):   
                    set.cluster = False    
                    setLabel = set.label    
                    setDepth = set.depth      
                    #determine the optimal number of clusters
                    print("Clustering set" + set.label)      
                    k = self.gapStatistic(set)    
                    if k > 1:     
                        #cluster and divide sets based on clustering 
                        #update the list of sets 
                        converged = False     
                        labels = KMeans(n_clusters = k).fit_predict(set.points)         
                        for i in range(k): 
                            ind = np.where(labels == i)[0]  
                            points = set.points[ind]
                            reg = Region(points, self.model, setLabel + str(i), depth=setDepth+1)         
                            reg.fitPCA()      
                            #append new region to new set  
                            newViableSets.append(reg)                                           
            #end of clustering                   
            viableSets.extend(newViableSets)     
        #end of while loop              
        
if __name__ == '__main__':

    
    filename = os.path.join("results_optimization", "")
    
    #model = BioProc(np.array(["protein_production", "protein_production", "protein_production", "protein_production", "protein_degradation", "protein_degradation", "Kd","hill", "protein_production", "protein_degradation", "Kd", "hill"]), model_mode=three_bit_processor_ext, parameter_values=param_values, avg_dev=30)                                         
    model = model_clb()
    solver = Solver(model)                         
    solver.run(filename, maxDepth=1) #do not cluster         