content: from __future__ import division
  from collections import defaultdict
  import numpy as np
  from base_sample import BaseSample
  from sklearn.cluster import AffinityPropagation as AP
  import pandas as pd
  from collections import Counter
  
  class StratifiedRandom(BaseSample):
      
      def __init__(self, data_frame, number_arms=2):
          super(StratifiedRandom, self).__init__(data_frame, number_arms)
  
      
      def create_stratum(self, column_names, **kwargs):
          '''
          Use affinity propagation to find number of strata for each column. 
          column_names is a list of the covariates to be split into strata and 
          used for classification. This funciton adds a column to the data frame
          for each column as column_name_strata that gives the strata designation
          for that variable.  The whole data frame is returned.
          '''
  
          for colname in column_names:
              X = self.data[colname].reshape(-1, 1)
              
              if np.isnan(X).any():
                  raise ValueError("There are NaN values in self.data[%s] that the \
                                    clustering algorithm can't handle" % colname)
                                    
              elif np.unique(self.data[colname]).shape[0] <=2:
                  string_name = colname+'_strata'
                  self.data[string_name] = self.data[colname].astype(int)
          
              else:
                  af_model = AP(damping = 0.9)
                  strata_groups = af_model.fit(X)
                  
                  #cluster_centers_indices = af.cluster_centers_indices_
                  #n_clusters_ = len(cluster_centers_indices)
                  
                  string_name = colname+'_strata'
                  self.data[string_name] = strata_groups.labels_
                  
          return self.data
          
      
      #In the main function, you need to call create_stratum before create_unique_strata       
      def create_unique_strata(self, column_names):
          '''
          The input should be self.data that has had the strata for each column
          name assigned and had a pre-seeded randomization, meaning each arm
          has at least one randomly assigned participant.
          '''
      
          #Create a column to store concatenated strata strings for each data point
          self.data['strata_string'] = np.ones(len(self.data))*np.nan    
          
          #Initialize variables to be filled in during the loop        
          strata_unique = {}     
          
          #Loop through data points and create their strata strings
          for ind in self.data.index.values:
              similar_val = ''
              for colname in column_names:
                  string_name = colname+'_strata'
                  similar_val += str(self.data[string_name].loc[ind])
    
              
              #Add the total strata string for that data point
              self.data['strata_string'].set_value(ind,similar_val)
              
              #If the strata string exists, continue. If not, assign it a new value
              if similar_val in list(strata_unique.keys()):
                  strata_unique[similar_val].append(ind)
                  continue
              else:
                  strata_unique[similar_val] = [ind]
          
          return (strata_unique, self.data)
              
         
      def count_arm_assignments(self, strata_unique, key):
          '''
          For each unique strata, count how many are assigned to each arm.
          '''
          #Initialize arm_tally that is the same length as the number of arms
          arm_tally = np.zeros(self.n_arms)              
          
          #Loop through the values in the unique strata and count how many are in each arm
          for value in strata_unique[key]:
              
              #If it is not NaN, add one to the arm_tally for the data point's arm assignment
              if np.isnan(self.data['arm_assignment'][value]) == False:
                  arm_tally[int(self.data['arm_assignment'][value]-1)] += 1;
          
          return arm_tally
  
  
             
      def assign_arms(self, column_names, percent_nan = 0.05):
          '''
          Loop through unique strata and assign each data point to an arm.
          '''
          #clear all values with NaNs
          self.data = self.nan_finder(column_names, percent_nan)        
          
          #call create_stratum to create strata for each chosen covariate 
          self.data = self.create_stratum(column_names,preference=-50)
          
          #combine the covariate strata into a unique strata identifier
          (strata_unique, self.data) = self.create_unique_strata(column_names)        
          
          #initiate an empty column in the data frame for arm assignments
          self.data['arm_assignment'] = np.ones(len(self.data))*np.nan   
  
          #Loop through the uniqie strata       
          for key in strata_unique.keys():
              #Loop through the values in the unique stratum
              for value in strata_unique[key]:
                  #update the arm_tally based on new assignments
                  arm_tally = self.count_arm_assignments(strata_unique, key);
                  
                  ind_unique = np.where(arm_tally==np.min(arm_tally))[0]
                  self.data['arm_assignment'].set_value(
                          value, np.random.choice(list(ind_unique+1)
                  ))
  
          return self.data
          
  
  #        
  