from importlibs import *


__author__  = "Andrea Maffezzoli"
__revision__ = "$Revision: 0.64 - 24 Feb 2011$"
__name__    = "burst_clusterizer"
__all__     = ["clusterize"]



class clusterize(object):
    """This class performs burst detection for each channel and
       calculates the related indexes of spiking activity"""
    
    _pca_on = pca_on
    processing_flow = []
    kmeans_dict = {'replicates_nr':  replicates_nr,
                   'scatter_method': kmeans.scatterwithin_no_samplesnr, 
                   'Mahalanobis_threshold': Mahalanobis_threshold,
                   'normalization':  False,
                   'scaling_down':   False}

    def __init__(self, bn, results_path = '', exp_name = '', exp_step = [], **csv_options):
        
        self.results_path, self.exp_step = results_path, exp_step
        self.processing_flow = [self.__preprocessing__, self.__cluster__]
        
        ##################################
        # When clustering just burst data
        ##################################
        try:
            self.bursts_mean = bn.bursts_mean_tocluster
            self.nr_vars     = len(self.bursts_mean.values()[0])
            self.ratio_feat_nrvars = self.nr_vars / len(bn.burst_features)
            if self.ratio_feat_nrvars == 1: self.features = bn.burst_features
            else:
                self.features = []
                for n in xrange(self.ratio_feat_nrvars):
                    self.features.extend([j + '_step' + self.exp_step[n] for j in bn.burst_features])
            self.neurons_id         = bn.neurons_id
            self.neuron_spiking     = bn.neuron_spiking
            self.neuron_not_spiking = bn.neuron_not_spiking
            self.desc_bursts_stat   = bn.desc_bursts_stat
        #################################
        # When clustering any other data
        #################################
        except:
            if isinstance(bn, dict):   self.bursts_mean = bn
            elif isinstance(bn, list): self.bursts_mean = dict(enumerate(bn))
            else: raise
            self.nr_vars    = len(self.bursts_mean.itervalues().next())
            self.features   = map(repr, xrange(self.nr_vars))
            self.neurons_id = self.neuron_spiking = self.bursts_mean.keys()
            self.neuron_not_spiking = []
        
        self.neurons_to_cluster = [neuron for neuron in self.neuron_spiking
                                   if None not in self.bursts_mean[neuron]]
        self.neurons_pruned_before_cluster = tuple(set(self.neuron_spiking) - set(self.neurons_to_cluster))
        
        #################################
        # Clustering settings
        #################################
        self.clustering = None
        self.kmeans_dict['data_name'] = self.neurons_to_cluster

    def __call__(self, pre_run = False, pre_run_cl_nr = 2, verbose = True):

        if pre_run: self.__preprocessing__(); self.__precluster__(pre_run_cl_nr)
        else: self.__preprocessing__(); self.__cluster__(verbose = verbose)
        if self.clustering:
            self.clusters_nr            = self.clustering.clusters_nr
            self.clusters_names         = self.clustering.clusters_names
            self.clusters_data          = dict([(clust, [self.bursts_mean[neuron] for neuron in cl_names]) \
                                                for clust, cl_names in self.clusters_names.iteritems()])
            self.clusters_centroids     = self.clustering.clusters_centroids
            self.centroids_plus_sem     = self.clustering.centroids_plus_sem
            self.samples_for_cluster    = self.clustering.samples_for_cluster
            self.nr_samples_for_cluster = self.clustering.nr_samples_for_cluster
            if self.Outliers_detect: self.all_outliers_names = self.clustering.all_outliers_names

            #######################################################
            # Giving results also in terms of descriptive features
            #######################################################
            try:
                self.clusters_desc_data = dict((clust, [self.desc_bursts_stat[neuron] for neuron in cl_names]) \
                                                for clust, cl_names in self.clusters_names.iteritems())
                self.centroids_desc_plus_sem = dict((clust, centroid_plus_sem(data)) \
                                                for clust, data in self.clusters_desc_data.iteritems())
            except:
                self.clusters_desc_data, self.centroids_desc_plus_sem = self.clusters_data, self.centroids_plus_sem

    def __preprocessing__(self):
        
        #############################################################
        # Performing clustering only for units with burst data, that
        # is, those with non empty elements of self.bursts_mean
        #############################################################
        D = [self.bursts_mean[neuron] for neuron in self.neurons_to_cluster]
        if not(D): raise Exception, "Input data matrix is empty!"
        
        #######################################
        # Checking how many features there are
        #######################################
        if self.nr_vars > 1:

            #############################################################################
            # CORRECTED: FALSE -> If multi step processing is on, data normalization is done by default
            #############################################################################
            #if self.ratio_feat_nrvars > 1:
            if auto_scale_Normalization:
                print 'Data normalization before PCA+k-means..'
                normalization, M = True, normalize(D)
            #####################################
            # CORRECTED: FALSE -> f data normalization is automatic
            #####################################
            #elif not(auto_scale_Normalization):
            else:
            #    print 'Auto normalization is off:'
            #    while True:
            #        normalization = raw_input("\nDo you want to normalize data?\n\
            #- press Enter if Yes,\n- digit N and then Enter if Not\n>> ")
            #        if normalization == '':
            #            normalization, M = True, normalize(D)
            #            print 'Normalizing data..'; break
            #        elif normalization == 'N':
                print 'No data normalization before PCA+k-means..'
                normalization, M = False, numpy.asarray(D)
            #########################################
            # If data normalization is not automatic
            #########################################
            #else:
            #    #print 'Auto normalization is on:',
            #    ####################################################################
            #    # Here check the original bn.burst_features. In the multi-step case
            #    # it should be necessary to handle the same features but with many 
            #    # different suffixes given by each experimental step involved 
            #    ####################################################################
            #    if set(self.features) <= already_normalized_features:
            #       normalization, M = False, copy.copy(D) ; print 'no normalizing data..'
            #    else: normalization, M = True, normalize(D); print 'normalizing data..'

            #########################################################
            # Storing covariances matrix of variables
            #########################################################
            if store_PCAdata:
                pca_results_path = self.results_path + 'PCA_data\/'
                if not(os.path.exists(pca_results_path)): os.mkdir(pca_results_path, 0777)    
                burstfeat_covmatrix = numpy.corrcoef(M, rowvar = 0, bias = 1) 
                csvname = pca_results_path + '-'.join(self.exp_step)
                csvwriterows_safe(csvname + '_CORREL_COEFF.csv', [[''] + self.features] + \
                                [[feat] + [fpformat.fix(q, 4) for q in p] \
                                for feat, p in zip(self.features, burstfeat_covmatrix)])

        else: self._pca_on, normalization, self.M = False, False, copy.copy(D)
        
        #########################################################################
        #####   N.B. If matrix M has None values, PCA algorithm won't work  #####
        #########################################################################
        if self._pca_on:
            P = pca.pca(M)
            if store_PCAdata:
                csvwriterows_safe(csvname + '_PCA_SCORE.csv', [['% tot.var.'] + \
                                ['PC ' + repr(i) for i in range(self.nr_vars)]] + list(P.SCORE))
                csvwriterows_safe(csvname + '_PCA_latent.csv', [[lat / sum(P.latent) for lat in P.latent]])
                #try: pcawriter_COEFF = csv.writer(open(csvname + '_PCA_COEFF.csv', 'wb'), **csv_options)               
                #except csv.Error, e: sys.exit('%s' % e) ; pass

            ######################################
            # Asking the number of PCs to analyze
            ######################################
            sum_of_variances = 0.
            for pc_nr, v in enumerate(P.latent):
                sum_of_variances += v
                if sum_of_variances / sum(P.latent) >= pca_min_sum_of_variances / 100.: break
            if 'Release' not in globals(): 
                for l in xrange(self.nr_vars):
                    print " PC%d gives %.2f%s of the total var. (cumul.: %.2f%s)" \
                            % (l + 1, 100. * P.latent[l] / sum(P.latent), chr(37), \
                            100. * sum(P.latent[:l + 1]) / sum(P.latent), chr(37))
                print 'The first %d PCs give %.3f%c of the total variance' % (pc_nr+1, pca_min_sum_of_variances, chr(37))
            
            ##########################
            # Storing data to cluster
            ##########################
            self.M = P.SCORE[:,:pc_nr+1].tolist() #[p.tolist()[:pc_nr] for p in P.SCORE]

            ###########################################
            # Storing PCAs coefficients
            ###########################################
            if store_PCAdata:
                all_pca_2_write = []
                for p in xrange(pc_nr):
                    pca_2_write = [repr(p + 1) + ' PC']
                    P_COEFF_transposed = map(None, *P.COEFF)
                    for i, j in zip(xrange(self.nr_vars), P_COEFF_transposed[p]):
                        pca_2_write += [self.features[i] + ':'] + [fpformat.fix(j, 4)]
                    all_pca_2_write.append(pca_2_write)
                csvwriterows_safe(csvname + '_PCA_COEFF.csv', all_pca_2_write)
                        
    def __precluster__(self, n = 2):
        
        print "Finding %d clusters..." % n
        kmr = kmeans.kmeans_replica(self.M, min_clusters_nr = n, max_clusters_nr = n,
                                    outliers_detect = True, **self.kmeans_dict)
        for i, nrs in enumerate(sorted(kmr.nr_samples_for_cluster, reverse = True), start = 1):
            print " Cluster %d contains %d neurons" % (i, nrs)
        self.Outliers_detect = True
        self.clustering = kmr 
        
    def __cluster__(self, verbose = True):
        
        while True:

            ############################
            # Asking to fix clusters nr and if to perform Outliers Detection on clusters
            ############################
            while True:
                try:
                    _clustersnr = raw_input("Digit the number of clusters you want: ")
                    if _clustersnr == '':
                        if 'Release' in globals(): minClustNr, maxClustNr = 3, 3
                        else: minClustNr, maxClustNr = 2, 2
                    else: minClustNr = maxClustNr = int(_clustersnr)
                    break
                except ValueError: continue
            while True:
                Outliers_detect = raw_input("To perform Outliers Detection press Enter, if not digit 'N' and then Enter: ")
                if Outliers_detect == '':  self.Outliers_detect = True;  break
                if Outliers_detect == 'N': self.Outliers_detect = False; break

            ##########################################################
            # Replicating clustering in order to have optimal results
            # Storing final clusters results in variables 
            # clusters_names is a dict having the ordinal of the cluster as key and 
            # as corresponding value the list of names of neurons in the gievn cluster.
            # clusters_data is a dict having the ordinal of the cluster as key and 
            # as corresponding value the list of variables values of neurons in the given cluster.
            ##########################################################
            print 'Performing clustering...'
            start = time.time()
            kmr = kmeans.kmeans_replica(self.M, min_clusters_nr = minClustNr, max_clusters_nr = maxClustNr,
                                        outliers_detect = self.Outliers_detect, **self.kmeans_dict)
            print "Clustering has been performed in %.3f seconds." % (time.time() - start)
            if verbose: 
                for i, nrs in enumerate(sorted(kmr.nr_samples_for_cluster, reverse = True)):
                    print 'Cluster %d contains %d neurons' % (i+1, nrs)
            
            if 'Release' not in globals(): break
            clustering_is_ok = raw_input("To repeat clustering digit 'Y' and then Enter, if not only Enter: ")
            if clustering_is_ok != 'Y': break
            del kmr
        self.clustering = kmr  
