from itertools import product
from collections import defaultdict

from AbstractRepairer import AbstractRepairer
from CategoricalFeature import CategoricalFeature
from calculators import get_median
from SparseList import SparseList

import random
import math
from copy import deepcopy

class Repairer(AbstractRepairer):
  def repair(self, data_to_repair):
    num_cols = len(data_to_repair[0])
    col_ids = range(num_cols)

    # Get column type information
    col_types = ["Y"]*len(col_ids)
    for i, col in enumerate(col_ids):
      if i in self.features_to_ignore:
        col_types[i] = "I"
      elif i == self.feature_to_repair:
        col_types[i] = "X"

    col_type_dict = {col_id: col_type for col_id, col_type in zip(col_ids, col_types)}

    not_I_col_ids = filter(lambda x: col_type_dict[x] != "I", col_ids)

    if self.kdd:
      cols_to_repair = filter(lambda x: col_type_dict[x] == "Y", col_ids)
    else:
      cols_to_repair = filter(lambda x: col_type_dict[x] in "YX", col_ids)

    # To prevent potential perils with user-provided column names, map them to safe column names
    safe_stratify_cols = [self.feature_to_repair]

    # Extract column values for each attribute in data
    # Begin byled code will usually be created in the same directory as the .py file. initializing keys and values in dictionary
    data_dict = {col_id: [] for col_id in col_ids}

    # Populate each attribute with its column values
    for row in data_to_repair:
      for i in col_ids:
        data_dict[i].append(row[i])

    repair_types = {}
    for col_id, values in data_dict.items():
      if all(isinstance(value, float) for value in values):
        repair_types[col_id] = float
      elif all(isinstance(value, int) for value in values):
        repair_types[col_id] = int
      else:
        repair_types[col_id] = str

    """
     Create unique value structures: When performing repairs, we choose median values. If repair is partial, then values will be modified to some intermediate value between the original and the median value. However, the partially repaired value will only be chosen out of values that exist in the data set.  This prevents choosing values that might not make any sense in the data's context.  To do this, for each column, we need to sort all unique values and create two data structures: a list of values, and a dict mapping values to their positions in that list. Example: There are unique_col_vals[col] = [1, 2, 5, 7, 10, 14, 20] in the column. A value 2 must be repaired to 14, but the user requests that data only be repaired by 50%. We do this by finding the value at the right index:
       index_lookup[col][2] = 1
       index_lookup[col][14] = 5
       this tells us that unique_col_vals[col][3] = 7 is 50% of the way from 2 to 14.
    """
    unique_col_vals = {}
    index_lookup = {}
    for col_id in not_I_col_ids:
      col_values = data_dict[col_id]
      # extract unique values from column and sort
      col_values = sorted(list(set(col_values)))
      unique_col_vals[col_id] = col_values
      # look up a value, get its position
      index_lookup[col_id] = {col_values[i]: i for i in range(len(col_values))}

    """
     Make a list of unique values per each stratified column.  Then make a list of combinations of stratified groups. Example: race and gender cols are stratified: [(white, female), (white, male), (black, female), (black, male)] The combinations are tuples because they can be hashed and used as dictionary keys.  From these, find the sizes of these groups.
    """
    unique_stratify_values = [unique_col_vals[i] for i in safe_stratify_cols]
    all_stratified_groups = list(product(*unique_stratify_values))
    # look up a stratified group, and get a list of indices corresponding to that group in the data
    stratified_group_indices = defaultdict(list)

    # Find the number of unique values for each strat-group, organized per column.
    val_sets = {group: {col_id:set() for col_id in cols_to_repair}
                                     for group in all_stratified_groups}
    for i, row in enumerate(data_to_repair):
      group = tuple(row[col] for col in safe_stratify_cols)
      for col_id in cols_to_repair:
        val_sets[group][col_id].add(row[col_id])

      # Also remember that this row pertains to this strat-group.
      stratified_group_indices[group].append(i)

    """
     Separate data by stratified group to perform repair on each Y column's values given that their corresponding protected attribute is a particular stratified group. We need to keep track of each Y column's values corresponding to each particular stratified group, as well as each value's index, so that when we repair the data, we can modify the correct value in the original data. Example: Supposing there is a Y column, "Score1", in which the 3rd and 5th scores, 70 and 90 respectively, belonged to black women, the data structure would look like: {("Black", "Woman"): {Score1: [(70,2),(90,4)]}}
    """
    stratified_group_data = {group: {} for group in all_stratified_groups}
    for group in all_stratified_groups:
      for col_id, col_dict in data_dict.items():
        # Get the indices at which each value occurs.
        indices = {}
        for i in stratified_group_indices[group]:
          value = col_dict[i]
          if value not in indices:
            indices[value] = []
          indices[value].append(i)

        stratified_col_values = [(occurs, val) for val, occurs in indices.items()]
        stratified_col_values.sort(key=lambda tup: tup[1])
        stratified_group_data[group][col_id] = stratified_col_values

    mode_feature_to_repair = get_mode(data_dict[self.feature_to_repair])

    # Repair Data and retrieve the results
    for col_id in cols_to_repair:
      # which bucket value we're repairing
      group_offsets = {group: 0 for group in all_stratified_groups}
      col = data_dict[col_id]

      num_quantiles = min(len(val_sets[group][col_id]) for group in all_stratified_groups)
      quantile_unit = 1.0/num_quantiles

      if repair_types[col_id] in {int, float}:
        for quantile in range(num_quantiles):
          median_at_quantiles = []
          indices_per_group = {}

          for group in all_stratified_groups:
            group_data_at_col = stratified_group_data[group][col_id]
            num_vals = len(group_data_at_col)
            offset = int(round(group_offsets[group]*num_vals))
            number_to_get = int(round((group_offsets[group] + quantile_unit)*num_vals) - offset)
            group_offsets[group] += quantile_unit

            if number_to_get > 0:

              # Get data at this quantile from this Y column such that stratified X = group
              offset_data = group_data_at_col[offset:offset+number_to_get]
              indices_per_group[group] = [i for val_indices, _ in offset_data for i in val_indices]
              values = sorted([float(val) for _, val in offset_data])

              # Find this group's median value at this quantile
              median_at_quantiles.append( get_median(values, self.kdd) )

          # Find the median value of all groups at this quantile (chosen from each group's medians)
          median = get_median(median_at_quantiles, self.kdd)
          median_val_pos = index_lookup[col_id][median]

          # Update values to repair the dataset.
          for group in all_stratified_groups:
            for index in indices_per_group[group]:
              original_value = col[index]

              current_val_pos = index_lookup[col_id][original_value]
              distance = median_val_pos - current_val_pos # distance between indices
              distance_to_repair = int(round(distance * self.repair_level))
              index_of_repair_value = current_val_pos + distance_to_repair
              repaired_value = unique_col_vals[col_id][index_of_repair_value]

              # Update data to repaired valued
              data_dict[col_id][index] = repaired_value

      #Categorical Repair is done below
      elif repair_types[col_id] in {str}:
        feature = CategoricalFeature(col)
        categories = feature.bin_index_dict.keys()

        group_features = get_group_data(all_stratified_groups, stratified_group_data, col_id)

        categories_count = get_categories_count(categories, all_stratified_groups, group_features)

        categories_count_norm = get_categories_count_norm(categories, all_stratified_groups, categories_count, group_features)

        median = get_median_per_category(categories, categories_count_norm)

        # Partially fill-out the generator functions to simplify later calls.
        dist_generator = lambda group_index, category : gen_desired_dist(group_index, category, col_id, median, self.repair_level, categories_count_norm, self.feature_to_repair, mode_feature_to_repair)

        count_generator = lambda group_index, group, category : gen_desired_count(group_index, group, category, median, group_features, self.repair_level, categories_count)

        group_features, overflow = flow_on_group_features(all_stratified_groups, group_features, count_generator)

        group_features, assigned_overflow, distribution = assign_overflow(all_stratified_groups, categories, overflow, group_features, dist_generator)

        # Return our repaired feature in the form of our original dataset
        for group in all_stratified_groups:
          indices = stratified_group_indices[group]
          for i, index in enumerate(indices):
            repaired_value = group_features[group].data[i]
            data_dict[col_id][index] = repaired_value

    # Replace stratified groups with their mode value, to remove it's information
    repaired_data = []
    for i, orig_row in enumerate(data_to_repair):
      new_row = [orig_row[j] if j not in cols_to_repair else data_dict[j][i] for j in col_ids]
      repaired_data.append(new_row)
    return repaired_data

def get_group_data(all_stratified_groups,stratified_group_data, col_id):
  group_features={}
  for group in all_stratified_groups:
    points = [(i, val) for indices, val in stratified_group_data[group][col_id] for i in indices]
    points = sorted(points, key=lambda x: x[0]) # Sort by index
    values = [value for _, value in points]

    # send values to CategoricalFeature object, which bins the data into categories
    group_features[group] = CategoricalFeature(values)
  return group_features

# Count the observations in each category. e.g. categories_count[1] = {'A':[1,2,3], 'B':[3,1,4]}, for column 1, category 'A' has 1 observation from group 'x', 2 from 'y', ect.
def get_categories_count(categories, all_stratified_groups, group_feature):
  # Grossness for speed efficiency. Don't worry, it make me sad, too.
  count_dict={cat: SparseList(data=(group_feature[group].category_count[cat] if cat in group_feature[group].category_count else 0 for group in all_stratified_groups)) for cat in categories}

  return count_dict

# Find the normalized count for each category, where normalized count is count divided by the number of people in that group
def get_categories_count_norm(categories, all_stratified_groups, count_dict, group_features):
  # Forgive me, Father, for I have sinned in bringing this monstrosity into being.
  norm = {cat: SparseList(data=(count_dict[cat][i] * (1.0/len(group_features[group].data)) if group_features[group].data else 0.0 for i,group in enumerate(all_stratified_groups))) for cat in categories}
  return norm

# Find the median normalized count for each category
def get_median_per_category(categories, categories_count_norm):
  return {cat: get_median(categories_count_norm[cat], False) for cat in categories}

# Generate the desired distribution and desired "count" for a given group-category-feature combination.
def gen_desired_dist(group_index, cat, col_id, median, repair_level, norm_counts, feature_to_remove, mode):
      if feature_to_remove == col_id:
        return 1 if cat==mode else (1-repair_level)*norm_counts[cat][group_index]
      else:
        return (1 - repair_level)*norm_counts[cat][group_index] + (repair_level*median[cat])

def gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count):
      med=median[category]
      size = len(group_features[group].data)
      # des-proportion = (1-lambda)*original-count  + (lambda)*median-count
      count = categories_count[category][group_index]
      des_count = math.floor(((1-repair_level)*count)+(repair_level)*med*size)
      return des_count

 # Run Max-flow to distribute as many observations to categories as possible. Overflow are those observations that are left over
def flow_on_group_features(all_stratified_groups, group_features, repair_generator):
  dict1= {}
  dict2={}
  for i, group in enumerate(all_stratified_groups):
    feature = group_features[group]
    count_generator = lambda category : repair_generator(i, group, category)

    # Create directed graph from nodes that supply the original counts to nodes that demand
    # the desired counts, with a overflow node as total desired count is at most total original counts
    DG=feature.create_graph(count_generator)

    # Run max-flow, and record overflow count (total and per-group)
    new_feature,overflow = feature.repair(DG)
    dict2[group] = overflow

    # Update our original values with the values from max-flow, Note: still missing overflowed observations
    dict1[group] = new_feature

  return dict1, dict2

# Assign overflow observations to categories based on the group's desired distribution
def assign_overflow(all_stratified_groups, categories, overflow, group_features, repair_generator):
  feature = deepcopy(group_features)
  assigned_overflow = {}
  desired_dict_list = {}
  for group_index, group in enumerate(all_stratified_groups):
    # Calculate the category proportions.
    dist_generator = lambda cat: repair_generator(group_index, cat)
    cat_props = map(dist_generator,categories)

    if all(elem==0 for elem in cat_props): #TODO: Check that this is correct!
      cat_props = [1.0/len(cat_props)] * len(cat_props)
    s = float(sum(cat_props))
    cat_props = [elem/s for elem in cat_props]
    desired_dict_list[group] = cat_props
    assigned_overflow[group] = {}
    for i in range(int(overflow[group])):
      distribution_list = desired_dict_list[group]
      number = random.uniform(0, 1)
      cat_index = 0
      tally = 0
      for j in range(len(distribution_list)):
        value=distribution_list[j]
        if number < (tally+value):
          cat_index = j
          break
        tally += value
      assigned_overflow[group][i] = categories[cat_index]
    # Actually do the assignment
    count = 0
    for i, value in enumerate(group_features[group].data):
      if value == 0:
        (feature[group].data)[i] = assigned_overflow[group][count]
        count += 1
  return feature, assigned_overflow, desired_dict_list

def get_mode(values):
  counts = {}
  for value in values:
    counts[value] = 1 if value not in counts else counts[value]+1
  mode_tuple = max(counts.items(), key=lambda tup: tup[1])
  return mode_tuple[0]

def test():
  test_minimal()
  test_get_group_data()
  test_get_categories_count()
  test_get_categories_count_norm()
  test_get_median_per_category()
  test_gen_desired_count()
  test_gen_desired_dist()
  test_assign_overflow()
  test_categorical()
  test_numerical()
  test_repeated_values()

def test_repeated_values():
  all_data = [
  ["x","A"], ["x","B"], ["x","C"], ["x","D"], ["x","E"],
  ["y","F"], ["y","G"], ["y","H"], ["y","I"],
  ["z","J"], ["z","K"], ["z","L"], ["z","M"], ["z","N"], ["z","O"]]

  random.seed(10)

  repair_level=.5
  feature_to_repair = 0
  repairer = Repairer(all_data, feature_to_repair, repair_level, False)
  repaired_data=repairer.repair(all_data)

  correct_repaired_data = [
  ['x', 'E'], ['x', 'C'], ['x', 'C'], ['z', 'E'], ['x', 'B'],
  ['y', 'I'], ['y', 'I'], ['y', 'G'], ['z', 'H'],
  ['z', 'N'], ['z', 'N'], ['z', 'L'], ['z', 'K'], ['z', 'K'], ['z', 'K']]

  print "Test unique values -- .5 repaired_data altered?", repaired_data != all_data
  print "Test unique values -- .5 repaired_data correct?", repaired_data == correct_repaired_data

  all_data = [
  ["x","A"], ["x","A"], ["x","A"], ["x","A"], ["x","A"],
  ["y","A"], ["y","A"], ["y","A"], ["y","A"],
  ["z","A"], ["z","A"], ["z","A"], ["z","A"], ["z","A"], ["z","A"]]

  random.seed(10)

  repair_level=.5
  feature_to_repair = 0
  repairer = Repairer(all_data, feature_to_repair, repair_level, False)
  repaired_data=repairer.repair(all_data)

  correct_repaired_data = [
  ['x', 'A'], ['x', 'A'], ['x', 'A'], ['z', 'A'], ['x', 'A'],
  ['y', 'A'], ['y', 'A'], ['y', 'A'], ['z', 'A'],
  ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A']]

  print "Test repeated values -- .5 repaired_data altered?", repaired_data != all_data
  print "Test repeated values -- .5 repaired_data correct?", repaired_data == correct_repaired_data

def test_minimal():
  class_1 = [[float(i),"A"] for i in xrange(0, 100)]
  class_2 = [[float(i),"B"] for i in xrange(101, 200)] # Thus, "A" is mode class.
  data = class_1 + class_2

  feature_to_repair = 1
  repairer = Repairer(data, feature_to_repair, 1, False)
  repaired_data = repairer.repair(data)
  print "Minimal Dataset -- repaired_data altered?", repaired_data != data

  mode = get_mode([row[feature_to_repair] for row in data])
  print "Minimal Dataset -- mode is true mode?", mode=="A"
  print "Minimal Dataset -- mode value as feature_to_repair?", all(row[feature_to_repair] == mode for row in repaired_data)

def test_get_group_data():
  group_features = {}
  group_size = {}
  col_id = 1
  all_stratified_groups = [('y',),('z',)]
  stratified_group_data = {('y',): {1: [([4, 7, 5], 'A'),([3, 2, 6], 'B'), ([], 'C')]},\
                           ('z',): {1: [([9, 1, 10], 'A'), ([], 'B'), ([11], 'C')]}}
  group_features[col_id] = get_group_data(all_stratified_groups, stratified_group_data, col_id)

  print "Test get_group_data -- group features correct?", \
    [group_features[col_id][group].data for group in all_stratified_groups] == [['B','B','A','A','B','A'],['A', 'A', 'A', 'C']]
  #print "Test get_group_data -- group sizes correct?", group_size[col_id] == {('y',): 6, ('z',):4}

def test_get_categories_count():
  categories_count = {}
  categories = {1:['A','B','C','D']}
  all_stratified_groups = [('y',),('z',)]
  col_id = 1
  group_features = {1:{('y',): CategoricalFeature(['C','A','C','B','A','C']),\
                      ('z',): CategoricalFeature(['B','B','D','D'])}}\

  categories_count[col_id] = get_categories_count(categories[col_id], all_stratified_groups, group_features[col_id])
  #SparseList makes us test for correctness strangely
  print "Test get_categories_count -- category counts correct?",\
   categories_count[col_id]['A'][0] ==2 and categories_count[col_id]['C'][0] ==3 and categories_count[col_id]['B'][0] ==1\
   and categories_count[col_id]['B'][1] ==2 and categories_count[col_id]['D'][1] ==2

def test_get_categories_count_norm():
  categories_count_norm = {}
  categories = {1:['A','B']}
  all_stratified_groups = [('y',),('z',)]
  col_id = 1
  categories_count = {1: {'A':[4,0],'B':[16,0]}}

  #group_size = {1: {('y',): 20,('z',): 0}}
  group_features = {1:{('y',): CategoricalFeature(['A','A','A','A','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B']),
                        ('z',): CategoricalFeature([])}}
  categories_count_norm[col_id] = get_categories_count_norm(categories[col_id], all_stratified_groups, categories_count[col_id], group_features[col_id])
  #SparseList makes us test for correctness strangely
  print "Test get_categories_count_norm -- normalized category counts correct?",\
    categories_count_norm[col_id]['A'][0] == 0.2 and  categories_count_norm[col_id]['B'][0] == 0.8

def test_get_median_per_category():
  categories = {1:['A','B','C','D']}
  col_id =1
  categories_count_norm = {1:{'A':[0.25,0.0],'C':[0.3,0.0],'B':[0.4,0.0],  'D':[0.6,0.4]}}
  median = get_median_per_category(categories[col_id], categories_count_norm[col_id])
  print "Test get_median_per_category -- medians are correct?", median == {'A':0.0,'C':0.0,'B':0.0,  'D':0.4}

def test_gen_desired_count():
  #Case 1: feature with category with no values
  group_index = 0
  group = ('y',)
  median = {'A': 0.0 ,'B':0.0}
  group_features =  {('y',): CategoricalFeature(['A','A','A','A','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B']),
                        ('z',): CategoricalFeature([])}
  repair_level = .25
  categories_count = {'A':[4,0],'B':[16,0]}
  category = 'B'
  des_countB = gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count)
  category = 'A'
  des_countA = gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count)
  print "Test gen_desired_count -- desired count correct for feature with category with no values?", des_countB==12 and des_countA==3

  #Case 2: feature with regular categories
  group_index = 0
  group = ('y',)
  median = {'A': 0.2 ,'B':0.75}
  group_features =  {('y',): CategoricalFeature(['A','A','A','A', 'A','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B']),
                        ('z',): CategoricalFeature(['A','A','B','B','B','B','B','B','B','B'])}
  repair_level = 0.001
  categories_count = {'A':[5,2],'B':[15,8]}
  category = 'B'
  des_countB = gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count)
  category = 'A'
  des_countA = gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count)
  print "Test gen_desired_count -- desired count correct for standard feature?", des_countB==15 and des_countA==4

  #Case 3: Repair feature after having repaired the other features on such feature
  col_id = 0
  median = {'Y': 0.0 ,'Z':0.0}
  group_features =  {('y',): CategoricalFeature(['Y','Y','Y','Y']),
                       ('z',): CategoricalFeature(['Z','Z'])}
  repair_level = .5
  categories_count = {'Y':[4,0],'Z':[0,2]}
  mode_feature = 'Y'
  group_index = 0
  group = ('y',)
  category = 'Y'
  des_countY = gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count)
  group_index = 1
  group = ('z',)
  category = 'Z'
  des_countZ = gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count)
  print "Test gen_desired_count -- desired count correct for mode category when repairing feature to remove?", des_countY==2 and des_countZ==1
  # If you are confused why desired count for Z is 2, it is beacuse our group_index is for group y

def test_gen_desired_dist():
  #Case 1: feature with category with no values
  group_index = 0
  #group = ('y',)
  col_id = 1
  median = {'A': 0.0 ,'B':0.0}
  #group_features =  {('y',): CategoricalFeature(['A','A','A','A','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B']),
  #                      ('z',): CategoricalFeature([])}
  repair_level = .5
  categories_count_norm =  {'A':[0.2,0.0],'B':[0.8,0.0]}
  #categories_count = {'A':[4,0],'B':[16,0]}
  feature_to_remove = 0
  mode_feature = 'B'
  category = 'B'
  des_distB = gen_desired_dist(group_index, category, col_id, median, repair_level, categories_count_norm, feature_to_remove, mode_feature)
  category = 'A'
  des_distA = gen_desired_dist(group_index, category, col_id, median, repair_level, categories_count_norm, feature_to_remove, mode_feature)
  print "Test gen_desired_dist -- desired distribution correct for feature with category with no values?", des_distB == .4 and des_distA == .1

  #Case 2: feature with regular categories
  group_index = 0
  #group = ('y',)
  median = {'A': 0.2 ,'B':0.75}
  #group_features =  {('y',): CategoricalFeature(['A','A','A','A', 'A','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B']),
  #                      ('z',): CategoricalFeature(['A','A','B','B','B','B','B','B','B','B'])}
  repair_level = 0.001
  categories_count_norm =  {'A':[0.25,0.2],'B':[0.75,0.8]}
  #categories_count = {'A':[5,2],'B':[15,8]}
  feature_to_remove = 0
  category = 'B'
  des_distB = gen_desired_dist(group_index, category, col_id, median, repair_level, categories_count_norm, feature_to_remove, mode_feature)
  category = 'A'
  des_distA = gen_desired_dist(group_index, category, col_id, median, repair_level, categories_count_norm, feature_to_remove, mode_feature)
  print "Test gen_desired_dist -- desired distribution correct for standard feature?", des_distB == .75 and des_distA == .24995

  #Case 3: Repair feature after having repaired the other features on such feature
  col_id = 0
  median = {'Y': 0.0 ,'Z':0.0}
  #group_features =  {('y',): CategoricalFeature(['Y','Y','Y','Y']),
  #                     ('z',): CategoricalFeature(['Z','Z'])}
  repair_level = .5
  categories_count_norm =  {'Y':[1.0,0.0],'Z':[0.0,1.0]}
  #categories_count = {'Y':[4,0],'Z':[0,2]}
  feature_to_remove = 0
  mode_feature = 'Y'
  group_index = 0
  category = 'Y'
  des_distY = gen_desired_dist(group_index, category, col_id, median, repair_level, categories_count_norm, feature_to_remove, mode_feature)
  group_index = 1
  category = 'Z'
  des_distZ = gen_desired_dist(group_index, category, col_id, median, repair_level, categories_count_norm, feature_to_remove, mode_feature)
  print "Test gen_desired_dist -- desired distribution correct for mode category when repairing feature to remove?", des_distY==1 and des_distZ==.5

def test_assign_overflow():
  group_index = 0
  group = ('y',)
  category = 'B'
  col_id = 1
  median = {'A': 0.0 ,'B':0.0}
  group_features =  {('y',): CategoricalFeature(['A','A','A','A','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B','B']),
                      ('z',): CategoricalFeature([])}
  repair_level = 1
  categories_count_norm =  {'A':[0.2,0.0],'B':[0.8,0.0]}
  categories_count = {'A':[4,0],'B':[16,0]}

  feature_to_remove = 0
  mode_feature = 'B'
  dist_generator = lambda group_index, category : gen_desired_dist(group_index, category, col_id, median, repair_level, categories_count_norm, feature_to_remove, mode_feature)
  count_generator = lambda group_index, group, category : gen_desired_count(group_index, group, category, median, group_features, repair_level, categories_count)

  random.seed(10)

  all_stratified_groups = [('y',),('z',)]
  categories = ['A','B']
  col_id = 1
  overflow = {('y',):2,('z',):2}
  group_features = {('y',):CategoricalFeature(['A','A','B','B',0,0]), ('z',): CategoricalFeature(['B',0,0])}

  feature, assigned_overflow, desired_dict_list = assign_overflow(all_stratified_groups, categories, overflow, group_features, dist_generator)
  print "Test assign_overflow -- updated group features correct?", \
   [feature[group].data for group in all_stratified_groups] ==[['A', 'A', 'B', 'B', 'B', 'A'], ['B', 'B', 'A']]
  print "Test assign_overflow -- assigned overflow correctly?", assigned_overflow == {('y',): {0: 'B', 1: 'A'}, ('z',): {0: 'B', 1: 'A'}}
  print "Test assign_overflow -- distribution correct?", desired_dict_list == {('y',): [0.5, 0.5], ('z',): [0.5, 0.5]}

def test_numerical():
  all_data = [ ["x", 1], ["x", 1], ["x", 2], ["x", 2], ["x", 3], ["x", 3],
               ["y", 1], ["y", 1], ["y", 2] ]
  feature_to_repair = 0
  repairer = Repairer(all_data, feature_to_repair, 1.0, False)
  repaired_data = repairer.repair(all_data)

  correct_repaired_data = [ ["x", 1], ["x", 1], ["x", 1], ["x", 1], ["x", 2], ["x", 2],
                            ["x", 1], ["x", 1], ["x", 2] ]

  print "Numerical minimal dataset -- full repaired_data correct?", repaired_data == correct_repaired_data

def test_categorical():
  all_data = [
  ["x","A"], ["x","A"], ["x","B"], ["x","B"], ["x","B"],
  ["y","A"], ["y","A"], ["y","A"], ["y","B"],
  ["z","A"], ["z","A"], ["z","A"], ["z","A"], ["z","A"], ["z","B"]]

  random.seed(10)

  repair_level=1
  feature_to_repair = 0
  repairer = Repairer(all_data, feature_to_repair, repair_level, False)
  repaired_data=repairer.repair(all_data)

  correct_repaired_data = [
  ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'B'],
  ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'B'],
  ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'B']]

  print "Categorical Minimal Dataset -- full repaired_data altered?", repaired_data != all_data
  print "Categorical Minimal Dataset -- full repaired_data correct?", repaired_data == correct_repaired_data

  repair_level=0.5
  feature_to_repair = 0
  repairer = Repairer(all_data, feature_to_repair, repair_level, False)
  part_repaired_data=repairer.repair(all_data)

  correct_part_repaired_data = [
  ['x', 'A'], ['x', 'A'], ['z', 'B'], ['z', 'B'], ['x', 'B'],
  ['y', 'A'], ['y', 'A'], ['y', 'A'], ['y', 'B'],
  ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'B']]

  print "Categorical Minimal Dataset -- partial (.5) repaired_data altered?", part_repaired_data != all_data
  print "Categorical Minimal Dataset -- partial (.5) repaired_data correct?", part_repaired_data == correct_part_repaired_data

  repair_level=0.2
  feature_to_repair = 0
  repairer = Repairer(all_data, feature_to_repair, repair_level, False)
  part_repaired_data=repairer.repair(all_data)

  correct_part2_repaired_data =  [
  ['x', 'A'], ['x', 'A'], ['x', 'B'], ['x', 'B'], ['z', 'B'],
  ['y', 'A'], ['y', 'A'], ['y', 'A'], ['y', 'B'],
  ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'A'], ['z', 'B']]

  print "Categorical Minimal Dataset -- partial (.2) repaired_data altered?", part_repaired_data != all_data
  print "Categorical Minimal Dataset -- partial (.2) repaired_data correct?", part_repaired_data == correct_part2_repaired_data

if __name__== "__main__":
  test()


import json
import os, shutil
import random
import csv

import matplotlib
matplotlib.use('Agg') # Set the back-end
import matplotlib.pyplot as plt

def load_audit_confusion_matrices(filename):
  """
  Loads a confusion matrix in a two-level dictionary format.

  For example, the confusion matrix of a 75%-accurate model
  that predicted 15 values (and mis-classified 5) may look like:
  {"A": {"A":10, "B": 5}, "B": {"B":5}}

  Note that raw boolean values are translated into strings, such that
  a value that was the boolean True will be returned as the string "True".
  """

  with open(filename) as audit_file:
    next(audit_file) # Skip the first line.

    # Extract the confusion matrices and repair levels from the audit file.
    confusion_matrices = []
    for line in audit_file:
      separator = ":"
      separator_index = line.index(separator)

      comma_index = line.index(',')
      repair_level = float(line[separator_index+2:comma_index])
      raw_confusion_matrix = line[comma_index+2:-2]
      confusion_matrix = json.loads( raw_confusion_matrix.replace("'","\"") )
      confusion_matrices.append( (repair_level, confusion_matrix) )

  # Sort the repair levels in case they are out of order for whatever reason.
  confusion_matrices.sort(key = lambda pair: pair[0])
  return confusion_matrices

def graph_audit(filename, measurers, output_image_file):
  with open(filename) as audit_file:
    header_line = audit_file.readline()[:-1] # Remove the trailing endline.

  confusion_matrices = load_audit_confusion_matrices(filename)

  x_axis = [repair_level for repair_level, _ in confusion_matrices]
  y_axes = []

  # Graph the results for each requested measurement.
  for measurer in measurers:
    y_axis = [measurer(matrix) for _, matrix in confusion_matrices]
    plt.plot(x_axis, y_axis, label=measurer.__name__)
    y_axes.append(y_axis)

  # Format and save the graph to an image file.
  plt.title(header_line)
  plt.axis([0,1,0,1.1]) # Make all the plots consistently sized.
  plt.xlabel("Repair Level")
  plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
  plt.savefig(output_image_file, bbox_inches='tight')
  plt.clf() # Clear the entire figure so future plots are empty.

  # Save the data used to generate that image file.
  with open(output_image_file + ".data", "w") as f:
    writer = csv.writer(f)
    headers = ["Repair Level"] + [calc.__name__ for calc in measurers]
    writer.writerow(headers)
    for i, repair_level in enumerate(x_axis):
      writer.writerow([repair_level] + [y_vals[i] for y_vals in y_axes])

def graph_audits(filenames, measurer, output_image_file):
  features = []
  y_axes = []
  for filename in filenames:
    with open(filename) as audit_file:
      header_line = audit_file.readline()[:-1] # Remove the trailing endline.
      feature = header_line[header_line.index(":")+1:]

    confusion_matrices = load_audit_confusion_matrices(filename)
    x_axis = [repair_level for repair_level, _ in confusion_matrices]
    y_axis = [measurer(matrix) for _, matrix in confusion_matrices]
    plt.plot(x_axis, y_axis, label=feature)

    features.append(feature)
    y_axes.append(y_axis)

  # Format and save the graph to an image file.
  plt.title(measurer.__name__)
  plt.axis([0,1,0,1.1]) # Make all the plots consistently sized.
  plt.xlabel("Repair Level")
  plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
  plt.savefig(output_image_file, bbox_inches='tight')
  plt.clf() # Clear the entire figure so future plots are empty.

  # Save the data used to generate that image file.
  with open(output_image_file + ".data", "w") as f:
    writer = csv.writer(f)
    headers = ["Repair Level"] + features
    writer.writerow(headers)
    for i, repair_level in enumerate(x_axis):
      writer.writerow([repair_level] + [y_vals[i] for y_vals in y_axes])

def rank_audit_files(filenames, measurer):
  scores = []
  for filename in filenames:
    with open(filename) as audit_file:
      header_line = audit_file.readline()[:-1] # Remove the trailing endline.
      feature = header_line[header_line.index(":")+1:]

    confusion_matrices = load_audit_confusion_matrices(filename)
    _, start_matrix = confusion_matrices[0]
    _, end_matrix = confusion_matrices[-1]
    score_difference = measurer(start_matrix)-measurer(end_matrix)
    scores.append( (feature, score_difference) )

  scores.sort(key = lambda score_tup: score_tup[1], reverse=True)
  return scores

def group_audit_ranks(filenames, measurer, similarity_bound=0.05):
  """
  Given a list of audit files, rank them using the `measurer` and
  return the features that never deviate more than `similarity_bound`
  across repairs.
  """

  def _partition_groups(feature_scores):
    groups = []
    for feature, score in feature_scores:
      added_to_group = False

      # Check to see if the feature belongs in a group with any other features.
      for i, group in enumerate(groups):
        mean_score, group_feature_scores = group
        if abs(mean_score - score) < similarity_bound:
          groups[i][1].append( (feature, score) )

          # Recalculate the representative mean.
          groups[i][0] = sum([s for _, s in group_feature_scores])/len(group_feature_scores)
          added_to_group = True
          break

      # If this feature did not much with the current groups, create another group.
      if not added_to_group:
        groups.append( [score, [(feature,score)]] )

    # Return just the features.
    return [[feature for feature, score in group] for _, group in groups]

  score_dict = {}
  features = []
  for filename in filenames:
    with open(filename) as audit_file:
      header_line = audit_file.readline()[:-1] # Remove the trailing endline.
      feature = header_line[header_line.index(":")+1:]
      features.append(feature)

    confusion_matrices = load_audit_confusion_matrices(filename)
    for rep_level, matrix in confusion_matrices:
      score = measurer(matrix)
      if rep_level not in score_dict:
        score_dict[rep_level] = {}
      score_dict[rep_level][feature] = score

  # Sort by repair level increasing repair level.
  score_keys = sorted(score_dict.keys())

  groups = [features]
  while score_keys:
    key = score_keys.pop()
    new_groups = []
    for group in groups:
      group_features = [(f, score_dict[key][f]) for f in group]
      sub_groups = _partition_groups(group_features)
      new_groups.extend(sub_groups)
    groups = new_groups

  return groups

def test():
  TMP_DIR = "tmp"
  if not os.path.exists(TMP_DIR):
    os.makedirs(TMP_DIR)

  test_contents = "GFA Audit for: Test Feature\n0.0.data:[0.0, {'A': {'B': 100}, 'B': {'B': 199}}]\n0.1.data:[0.1, {'A': {'B': 100}, 'B': {'B': 199}}]\n0.5.data:[0.5, {'A': {'B': 100}, 'B': {'B': 199}}]\n1.0.data:[1.0, {'A': {'B': 100}, 'B': {'B': 199}}]\n"
  test_filenames = [TMP_DIR + "/test_audit_1.audit",
                    TMP_DIR + "/test_audit_2.audit"]

  # Prepare the sample audit files.
  for filename in test_filenames:
    with open(filename, "w") as f:
      f.write(test_contents)

  # A mock measurement measurer that returns a random number.
  def mock_measurer(conf_matrix):
    return random.random()

  # Perform the audit and save it an output image.
  measurers = [mock_measurer, mock_measurer]
  output_image = TMP_DIR + "/test_image.png"
  graph_audit(test_filenames[0], measurers, output_image) # Only need to test 1.

  file_not_empty = os.path.getsize(output_image) > 0
  print("image file generated? --", file_not_empty)

  file_not_empty = os.path.getsize(output_image + ".data") > 0
  print("data file generated? --", file_not_empty)

  ranked_features = rank_audit_files(test_filenames, mock_measurer)
  print("ranked features sorted? --", ranked_features[0] > ranked_features[1])

  output_image = TMP_DIR + "/test_image2.png"
  graph_audits(test_filenames, mock_measurer, output_image)
  file_not_empty = os.path.getsize(output_image) > 0
  print("ranked image file generated? --", file_not_empty)

  shutil.rmtree(TMP_DIR)

if __name__=="__main__":
  test()


def accuracy(conf_matrix):
  """
  Given a confusion matrix, returns the accuracy.
  Accuracy Definition: http://research.ics.aalto.fi/events/eyechallenge2005/evaluation.shtml
  """
  total, correct = 0.0, 0.0
  for true_response, guess_dict in conf_matrix.items():
    for guess, count in guess_dict.items():
      if true_response == guess:
        correct += count
      total += count
  return correct/total

def BCR(conf_matrix):
  """
  Given a confusion matrix, returns Balanced Classification Rate.
  BCR is (1 - Balanced Error Rate).
  BER Definition: http://research.ics.aalto.fi/events/eyechallenge2005/evaluation.shtml
  """
  parts = []
  for true_response, guess_dict in conf_matrix.items():
    error = 0.0
    total = 0.0
    for guess, count in guess_dict.items():
      if true_response != guess:
        error += count
      total += count
    parts.append(error/total)
  BER = sum(parts)/len(parts)
  return 1 - BER

def get_conf_matrix(prediction_tuples):
  # Produce a confusion matrix in a dictionary format from those predictions.
  conf_table = {}
  for actual, guess in prediction_tuples:
    guess = convert_to_type(actual, guess) # ... Since file-reading changes types.

    if not actual in conf_table:
      conf_table[actual] = {}

    if not guess in conf_table[actual]:
      conf_table[actual][guess] = 1
    else:
      conf_table[actual][guess] += 1

  return conf_table

def convert_to_type(actual, guess):
  if type(actual) == bool:
    guess = guess=="True"
  else:
    actual_type = type(actual)
    guess = actual_type(guess)
  return guess

def test():
  test_conf_matrix()
  test_measurers()

def test_conf_matrix():
  pred_tuples = [(1,1),(1,1),(2,2),(3,3),(1,3),(3,1)]
  correct_conf_matrix = {1:{1:2, 3:1}, 2:{2:1}, 3:{1:1, 3:1}}
  conf_matrix = get_conf_matrix(pred_tuples)
  print "confusion matrix correct? -- ", conf_matrix == correct_conf_matrix

def test_measurers():
  conf_matrix = {"A":{"A":10}, "B":{"A":5,"B":5}}
  print "measurements -- accuracy correct? -- ", accuracy(conf_matrix) == 0.75
  print "measurements -- 1-BER correct? -- ", BCR(conf_matrix) == 0.75

  conf_matrix = {"A":{"A":10}}
  print "measurements -- accuracy correct? -- ", accuracy(conf_matrix) == 1.0
  print "measurements -- 1-BER correct? -- ", BCR(conf_matrix) == 1.0

  conf_matrix = {"A":{"A":95}, "B":{"A":5}}
  print "measurements -- accuracy correct? -- ", accuracy(conf_matrix) == 0.95
  print "measurements -- 1-BER correct? -- ", BCR(conf_matrix) == 0.50

if __name__=="__main__":
  test()


from copy import deepcopy

def get_median(values, kdd):
  """
  Given an unsorted list of numeric values, return median value (as a float).
  Note that in the case of even-length lists of values, we apply the value to
  the left of the center to be the median (such that the median can only be
  a value from the list of values).
  Eg: get_median([1,2,3,4]) == 2, not 2.5.
  """

  if not values:
    raise Exception("Cannot calculate median of list with no values!")

  sorted_values = deepcopy(values)
  sorted_values.sort() # Not calling `sorted` b/c `sorted_values` may not be list.

  if kdd:
    return sorted_values[len(values)//2]
  else:
    if len(values) % 2 == 0:
      return sorted_values[len(values)//2-1]
    else:
      return sorted_values[len(values)//2]

def test():
  test_median()

def test_median():
  feature_values = [4,1,3,2]
  correct_median = 2
  # not KDD repair
  print("Regular median value is correct?", get_median(feature_values, False) == correct_median)
  # KDD repair
  correct_median = 3
  print("KDD median value is correct?", get_median(feature_values, True) == correct_median)

if __name__=="__main__":
  test()


from AbstractModelFactory import AbstractModelFactory
from AbstractModelVisitor import AbstractModelVisitor

import os
import numpy as np
import tensorflow as tf
import datetime

#TMP_DIR = "tmp/"
#CHECKPOINT_DIR = "tmp/tensorflow_checkpoints/"
#for directory in [TMP_DIR, CHECKPOINT_DIR]:
#  if not os.path.exists(directory):
#    os.makedirs(directory)

class ModelFactory(AbstractModelFactory):

  def __init__(self, *args, **kwargs):
    super(ModelFactory, self).__init__(*args, **kwargs)
    self.verbose_factory_name = "TensorFlow_Network"
    self.model_name = "Recidivism TensorFlow"

    self.num_epochs = 2000
    self.batch_size = 500

    self.response_index = self.headers.index(self.response_header)

    possible_values = set(row[self.response_index] for row in self.all_data)
    self.num_labels = len(possible_values)
    self.response_dict = {val:i for i,val in enumerate(possible_values)}

    self.hidden_layer_sizes = [100, 40] # If empty, no hidden layers are used.
    self.layer_types = [tf.nn.softmax,
			tf.nn.softmax,   # Input Layer
                        tf.nn.softmax]     # 4th Hidden Layer

  def build(self, train_set): #TODO: Add a features-to-ignore option.
    # In case the class is a string, translate it.
    translated_train_set = translate_response(self.response_index, train_set, self.response_dict)

    train_matrix, train_labels = list_to_tf_input(translated_train_set, self.response_index, self.num_labels)
    train_size, num_features = train_matrix.shape

    # Construct the layer architecture.
    x = tf.placeholder("float", shape=[None, num_features]) # Input
    y_ = tf.placeholder("float", shape=[None, self.num_labels]) # Output.

    layer_sizes = [num_features] + self.hidden_layer_sizes + [self.num_labels]
    # Generate a layer for the input and for each additional hidden layer.
    layers = [x] # Count the input as the first layer.
    for i in xrange(len(layer_sizes)-1):
      layer_size = layer_sizes[i]
      layer_type = self.layer_types[i]

      prev_layer = layers[-1]
      next_layer_size = layer_sizes[i+1]

      # Create a new layer with initially random weights and biases.
      W = tf.Variable(tf.random_normal([layer_size, next_layer_size]))
      b = tf.Variable(tf.random_normal([next_layer_size]))
      new_layer = layer_type(tf.matmul(prev_layer, W) + b)

      layers.append( new_layer )

    y = layers[-1]

    # Optimization.
    cross_entropy = -tf.reduce_sum(y_*tf.log(y))
    train_step = tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy)

    saver = tf.train.Saver()  # Defaults to saving all variables.

    # Create a local session to run this computation.
    with tf.Session() as tf_session:
      # For the test data, hold the entire dataset in one constant node.
      tf.initialize_all_variables().run()

      # Iterate and train.
      for step in xrange(self.num_epochs * train_size // self.batch_size):

        offset = (step * self.batch_size) % train_size
        batch_data = train_matrix[offset:(offset + self.batch_size), :]
        batch_labels = train_labels[offset:(offset + self.batch_size)]
        train_step.run(feed_dict={x: batch_data, y_: batch_labels})

        # Save the model file each step.
        saver.save(tf_session, CHECKPOINT_DIR + 'model.ckpt', global_step=step+1)

    return ModelVisitor(saver, self.response_index, self.num_labels, x, y_, y, self.response_dict)

class ModelVisitor(AbstractModelVisitor):

  def __init__(self, model_saver, response_index, num_labels, x, y_, y, response_dict):
    self.model_saver = model_saver
    self.model_name = "Recidivism TensorFlow Visitor"
    self.response_index = response_index
    self.num_labels = num_labels
    self.x = x
    self.y_ = y_
    self.y = y
    self.response_dict = response_dict

  def test(self, test_set, test_name=""):
    translated_test_set = translate_response(self.response_index, test_set, self.response_dict)

    test_matrix, test_labels = list_to_tf_input(translated_test_set, self.response_index, self.num_labels  )

    with tf.Session() as tf_session:
      ckpt = tf.train.get_checkpoint_state(CHECKPOINT_DIR)
      self.model_saver.restore(tf_session, ckpt.model_checkpoint_path)
      predictions = tf.argmax(self.y, 1).eval(feed_dict={self.x: test_matrix, self.y_:test_labels}, session=tf_session)

    predictions_dict = {i:key for key,i in self.response_dict.items()}
    predictions = [predictions_dict[pred] for pred in predictions]

    return zip([row[self.response_index] for row in test_set], predictions)

def list_to_tf_input(data, response_index, num_labels):
  data = expand_to_one_hot(data)
  matrix = np.matrix([row[:response_index] + row[response_index+1:] for row in data])

  labels = np.asarray([row[response_index] for row in data], dtype=np.uint8)
  labels_onehot = (np.arange(num_labels) == labels[:, None]).astype(np.float32)

  return matrix, labels_onehot

def translate_response(response_index, data_set, response_dict):
  translated_set = []
  for row in data_set:
    response = row[response_index]
    translation = response_dict[response]
    new_row = row[:response_index]+[translation]+row[response_index+1:]
    translated_set.append(new_row)
  return translated_set

def expand_to_one_hot(data,expand = True,use_alternative=False):
    header_dict = {'ALCABUS':0,'PRIRCAT':1,'TMSRVC':2,'SEX1':3,'RACE':4,'RELTYP':5,'age_1st_arrest':6,'DRUGAB':7,'Class':8,'RLAGE':9,'NFRCTNS':10}

    new_data = []
    for entry in data:
	temp = {}
	if expand == True:
	    if entry[header_dict["SEX1"]] == "FEMALE":
		temp['female'] = 1
	    else:
		temp['female'] = 0

	    if entry[header_dict["ALCABUS"]] == 'INMATE IS AN ALCOHOL ABUSER':
		temp['prior_alcohol_abuse'] = 1
	    else:
		temp['prior_alcohol_abuse'] = 0

	    if entry[header_dict['DRUGAB']] == 'INMATE IS A DRUG ABUSER':
		temp['prior_drug_abuse'] = 1
	    else:
		temp['prior_drug_abuse'] = 0

	    if entry[header_dict['NFRCTNS']] == 'INMATE HAS RECORD':
		temp['infraction_in_prison'] = 1
	    else:
		temp['infraction_in_prison'] = 0

	    race_cats = ['WHITE','BLACK','AMERICAN INDIAN/ALEUTIAN','ASIAN/PACIFIC ISLANDER','OTHER','UNKNOWN']

	    for cat in race_cats:
		if entry[header_dict['RACE']] == cat:
		    temp['race_'+cat] = 1
		else:
		    temp['race_'+cat] = 0

	    release_age_cats = ['14 TO 17 YEARS OLD','18 TO 24 YEARS OLD', '25 TO 29 YEARS OLD', \
	    '30 TO 34 YEARS OLD','35 TO 39 YEARS OLD','40 TO 44 YEARS OLD','45 YEARS OLD AND OLDER']
	    for cat in release_age_cats:
		if entry[header_dict['RLAGE']] == cat:
		    temp['release_age_'+cat] = 1
		else:
		    temp['release_age_'+cat] = 0

	    time_served_cats = ['None','1 TO 6 MONTHS','13 TO 18 MONTHS','19 TO 24 MONTHS','25 TO 30 MONTHS', \
			'31 TO 36 MONTHS','37 TO 60 MONTHS','61 MONTHS AND HIGHER','7 TO 12 MONTHS']
	    for cat in time_served_cats:
		if entry[header_dict['TMSRVC']] == cat:
		    temp['time_served_'+cat] = 1
		else:
		    temp['time_served_'+cat] = 0

	    prior_arrest_cats = ['None','1 PRIOR ARREST','11 TO 15 PRIOR ARRESTS','16 TO HI PRIOR ARRESTS','2 PRIOR ARRESTS', \
		'3 PRIOR ARRESTS','4 PRIOR ARRESTS','5 PRIOR ARRESTS','6 PRIOR ARRESTS','7 TO 10 PRIOR ARRESTS']
	    for cat in prior_arrest_cats:
		if entry[header_dict['PRIRCAT']] == cat:
		    temp['prior_arrest_'+cat] = 1
		else:
		    temp['prior_arrest_'+cat] = 0

	    conditional_release =['PAROLE BOARD DECISION-SERVED NO MINIMUM','MANDATORY PAROLE RELEASE', 'PROBATION RELEASE-SHOCK PROBATION', \
			'OTHER CONDITIONAL RELEASE']
	    unconditional_release = ['EXPIRATION OF SENTENCE','COMMUTATION-PARDON','RELEASE TO CUSTODY, DETAINER, OR WARRANT', \
			'OTHER UNCONDITIONAL RELEASE']
	    other_release = ['NATURAL CAUSES','SUICIDE','HOMICIDE BY ANOTHER INMATE','OTHER HOMICIDE','EXECUTION','OTHER TYPE OF DEATH', \
		    'TRANSFER','RELEASE ON APPEAL OR BOND','OTHER TYPE OF RELEASE','ESCAPE','ACCIDENTAL INJURY TO SELF','UNKNOWN']
	    if entry[header_dict['RELTYP']] in conditional_release:
		temp['released_conditional'] = 1
		temp['released_unconditional'] = 0
		temp['released_other'] = 0
	    elif entry[header_dict['RELTYP']] in unconditional_release:
		temp['released_conditional'] = 0
		temp['released_unconditional'] = 1
		temp['released_other'] = 0
	    else:
		temp['released_conditional'] = 0
		temp['released_unconditional'] = 0
		temp['released_other'] = 1

	    first_arrest_cats = ['UNDER 17','BETWEEN 18 AND 24','BETWEEN 25 AND 29','BETWEEN 30 AND 39','OVER 40']
	    for cat in first_arrest_cats:
		if entry[header_dict['age_1st_arrest']] == cat:
		    temp['age_first_arrest_'+cat] = 1
		else:
		    temp['age_first_arrest_'+cat] = 0
	else:
	    temp['SEX1'] = entry['SEX1']
	    temp['RELTYP'] = entry['RELTYP']
	    temp['PRIRCAT'] = entry['PRIRCAT']
	    temp['ALCABUS'] = entry['ALCABUS']
	    temp['DRUGAB'] = entry['DRUGAB']
	    temp['RLAGE'] = entry['RLAGE']
	    temp['TMSRVC'] = entry['TMSRVC']
	    temp['NFRCTNS'] = entry['NFRCTNS']
	    temp['RACE'] = entry['RACE']
	    try:
		bdate = datetime.date(int(entry['YEAROB2']),int(entry['MNTHOB2']), int(entry['DAYOB2']))
		first_arrest = datetime.date(int(entry['A001YR']),int(entry['A001MO']),int(entry['A001DA']))
		first_arrest_age = first_arrest - bdate
		temp['age_1st_arrest'] = first_arrest_age.days
	    except:
		temp['age_1st_arrest'] = 0

	new_data.append(temp)

    # convert from dictionary to list of lists
    fin = [[int(entry[key]) for key in entry.keys()] for entry in new_data]
    """
    with open("brandon_testing/test_"+str(time.clock())+".csv","w") as f:
	writer = csv.writer(f,delimiter=",")
	for row in fin:
	    writer.writerow(row)
    """

    return fin

def test():
  # expand_to_one_hot assumes the data is the recidivism data, so these tests can't be run
  # test_list_to_tf_input()
  # test_basic_model()
  # test_categorical_model()
  pass

def test_list_to_tf_input():
  data = [[0,0],[0,1],[0,2]]
  tf_matrix, tf_onehot = list_to_tf_input(data, 1, 3)
  correct_matrix = [[0],[0],[0]]
  correct_onehot = [[1,0,0], [0,1,0], [0,0,1]]
  print "list_to_tf_input matrix correct? --",np.array_equal(tf_matrix, correct_matrix)
  print "list_to_tf_input onehot correct? --",np.array_equal(tf_onehot, correct_onehot)

def test_basic_model():
  headers = ["predictor 1", "predictor 2", "response"]
  response = "response"
  train_set = [[i,0,0] for i in range(1,50)] + [[0,i,1] for i in range(1,50)]
  test_set = [[i,0,0] for i in range(1,50)] + [[0,i,1] for i in range(1,50)]
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")
  print "factory settings valid? -- ",len(factory.hidden_layer_sizes)+1 == len(factory.layer_types)

  model = factory.build(train_set)
  print "factory builds ModelVisitor? -- ", isinstance(model, ModelVisitor)

  predictions = model.test(test_set)
  resp_index = headers.index(response)
  intended_predictions = [(row[resp_index], row[resp_index]) for row in test_set]
  print "predicting numeric categories correctly? -- ", predictions == intended_predictions

def test_categorical_model():
  headers = ["predictor 1", "predictor 2", "response"]
  response = "response"
  train_set = [[i,0,"A"] for i in range(1,50)] + [[0,i,"B"] for i in range(1,50)]
  test_set = [[i,0,"A"] for i in range(1,50)] + [[0,i,"C"] for i in range(1,50)]
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")
  print "factory settings valid? -- ",len(factory.hidden_layer_sizes)+1 == len(factory.layer_types)

  model = factory.build(train_set)
  print "factory builds ModelVisitor? -- ", isinstance(model, ModelVisitor)

  predictions = model.test(test_set)
  resp_index = headers.index(response)
  intended_predictions = [(test_row[resp_index], train_row[resp_index]) for train_row, test_row in zip(train_set,test_set)]
  print "predicting string-categories correctly? -- ", predictions == intended_predictions

if __name__=="__main__":
  test()


import json
import os, shutil
import random
import csv

import matplotlib
matplotlib.use('Agg') # Set the back-end
import matplotlib.pyplot as plt

def load_audit_confusion_matrices(filename):
  """
  Loads a confusion matrix in a two-level dictionary format.

  For example, the confusion matrix of a 75%-accurate model
  that predicted 15 values (and mis-classified 5) may look like:
  {"A": {"A":10, "B": 5}, "B": {"B":5}}

  Note that raw boolean values are translated into strings, such that
  a value that was the boolean True will be returned as the string "True".
  """

  with open(filename) as audit_file:
    audit_file.next() # Skip the first line.

    # Extract the confusion matrices and repair levels from the audit file.
    confusion_matrices = []
    for line in audit_file:
      separator = ":"
      separator_index = line.index(separator)

      comma_index = line.index(',')
      repair_level = float(line[separator_index+2:comma_index])
      raw_confusion_matrix = line[comma_index+2:-2]
      confusion_matrix = json.loads( raw_confusion_matrix.replace("'","\"") )
      confusion_matrices.append( (repair_level, confusion_matrix) )

  # Sort the repair levels in case they are out of order for whatever reason.
  confusion_matrices.sort(key = lambda pair: pair[0])
  return confusion_matrices

def graph_audit(filename, measurers, output_image_file):
  with open(filename) as audit_file:
    header_line = audit_file.readline()[:-1] # Remove the trailing endline.

  confusion_matrices = load_audit_confusion_matrices(filename)

  x_axis = [repair_level for repair_level, _ in confusion_matrices]
  y_axes = []

  # Graph the results for each requested measurement.
  for measurer in measurers:
    y_axis = [measurer(matrix) for _, matrix in confusion_matrices]
    plt.plot(x_axis, y_axis, label=measurer.__name__)
    y_axes.append(y_axis)

  # Format and save the graph to an image file.
  plt.title(header_line)
  plt.axis([0,1,0,1.1]) # Make all the plots consistently sized.
  plt.xlabel("Repair Level")
  plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
  plt.savefig(output_image_file, bbox_inches='tight')
  plt.clf() # Clear the entire figure so future plots are empty.

  # Save the data used to generate that image file.
  with open(output_image_file + ".data", "w") as f:
    writer = csv.writer(f)
    headers = ["Repair Level"] + [calc.__name__ for calc in measurers]
    writer.writerow(headers)
    for i, repair_level in enumerate(x_axis):
      writer.writerow([repair_level] + [y_vals[i] for y_vals in y_axes])

def graph_audits(filenames, measurer, output_image_file):
  features = []
  y_axes = []
  for filename in filenames:
    with open(filename) as audit_file:
      header_line = audit_file.readline()[:-1] # Remove the trailing endline.
      feature = header_line[header_line.index(":")+1:]

    confusion_matrices = load_audit_confusion_matrices(filename)
    x_axis = [repair_level for repair_level, _ in confusion_matrices]
    y_axis = [measurer(matrix) for _, matrix in confusion_matrices]
    plt.plot(x_axis, y_axis, label=feature)

    features.append(feature)
    y_axes.append(y_axis)

  # Format and save the graph to an image file.
  plt.title(measurer.__name__)
  plt.axis([0,1,0,1.1]) # Make all the plots consistently sized.
  plt.xlabel("Repair Level")
  plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
  plt.savefig(output_image_file, bbox_inches='tight')
  plt.clf() # Clear the entire figure so future plots are empty.

  # Save the data used to generate that image file.
  with open(output_image_file + ".data", "w") as f:
    writer = csv.writer(f)
    headers = ["Repair Level"] + features
    writer.writerow(headers)
    for i, repair_level in enumerate(x_axis):
      writer.writerow([repair_level] + [y_vals[i] for y_vals in y_axes])

def rank_audit_files(filenames, measurer):
  scores = []
  for filename in filenames:
    with open(filename) as audit_file:
      header_line = audit_file.readline()[:-1] # Remove the trailing endline.
      feature = header_line[header_line.index(":")+1:]

    confusion_matrices = load_audit_confusion_matrices(filename)
    _, start_matrix = confusion_matrices[0]
    _, end_matrix = confusion_matrices[-1]
    score_difference = measurer(start_matrix)-measurer(end_matrix)
    scores.append( (feature, score_difference) )

  scores.sort(key = lambda score_tup: score_tup[1], reverse=True)
  return scores

def group_audit_ranks(filenames, measurer, similarity_bound=0.05):
  """
  Given a list of audit files, rank them using the `measurer` and
  return the features that never deviate more than `similarity_bound`
  across repairs.
  """

  def _partition_groups(feature_scores):
    groups = []
    for feature, score in feature_scores:
      added_to_group = False

      # Check to see if the feature belongs in a group with any other features.
      for i, group in enumerate(groups):
        mean_score, group_feature_scores = group
        if abs(mean_score - score) < similarity_bound:
          groups[i][1].append( (feature, score) )

          # Recalculate the representative mean.
          groups[i][0] = sum([s for _, s in group_feature_scores])/len(group_feature_scores)
          added_to_group = True
          break

      # If this feature did not much with the current groups, create another group.
      if not added_to_group:
        groups.append( [score, [(feature,score)]] )

    # Return just the features.
    return [[feature for feature, score in group] for _, group in groups]

  score_dict = {}
  features = []
  for filename in filenames:
    with open(filename) as audit_file:
      header_line = audit_file.readline()[:-1] # Remove the trailing endline.
      feature = header_line[header_line.index(":")+1:]
      features.append(feature)

    confusion_matrices = load_audit_confusion_matrices(filename)
    for rep_level, matrix in confusion_matrices:
      score = measurer(matrix)
      if rep_level not in score_dict:
        score_dict[rep_level] = {}
      score_dict[rep_level][feature] = score

  # Sort by repair level increasing repair level.
  score_keys = sorted(score_dict.keys())

  groups = [features]
  while score_keys:
    key = score_keys.pop()
    new_groups = []
    for group in groups:
      group_features = [(f, score_dict[key][f]) for f in group]
      sub_groups = _partition_groups(group_features)
      new_groups.extend(sub_groups)
    groups = new_groups

  return groups

def test():
  TMP_DIR = "tmp"
  if not os.path.exists(TMP_DIR):
    os.makedirs(TMP_DIR)

  test_contents = "GFA Audit for: Test Feature\n0.0.data:[0.0, {'A': {'B': 100}, 'B': {'B': 199}}]\n0.1.data:[0.1, {'A': {'B': 100}, 'B': {'B': 199}}]\n0.5.data:[0.5, {'A': {'B': 100}, 'B': {'B': 199}}]\n1.0.data:[1.0, {'A': {'B': 100}, 'B': {'B': 199}}]\n"
  test_filenames = [TMP_DIR + "/test_audit_1.audit",
                    TMP_DIR + "/test_audit_2.audit"]

  # Prepare the sample audit files.
  for filename in test_filenames:
    with open(filename, "w") as f:
      f.write(test_contents)

  # A mock measurement measurer that returns a random number.
  def mock_measurer(conf_matrix):
    return random.random()

  # Perform the audit and save it an output image.
  measurers = [mock_measurer, mock_measurer]
  output_image = TMP_DIR + "/test_image.png"
  graph_audit(test_filenames[0], measurers, output_image) # Only need to test 1.

  file_not_empty = os.path.getsize(output_image) > 0
  print "image file generated? --", file_not_empty

  file_not_empty = os.path.getsize(output_image + ".data") > 0
  print "data file generated? --", file_not_empty

  ranked_features = rank_audit_files(test_filenames, mock_measurer)
  print "ranked features sorted? --", ranked_features[0] > ranked_features[1]

  output_image = TMP_DIR + "/test_image2.png"
  graph_audits(test_filenames, mock_measurer, output_image)
  file_not_empty = os.path.getsize(output_image) > 0
  print "ranked image file generated? --", file_not_empty

  shutil.rmtree(TMP_DIR)

if __name__=="__main__":
  test()


from BlackBoxAuditing.model_factories.AbstractModelFactory import AbstractModelFactory
from BlackBoxAuditing.model_factories.AbstractModelVisitor import AbstractModelVisitor

from sklearn import svm
import numpy as np
import copy

class ModelFactory(AbstractModelFactory):
    def __init__(self, *args, **kwargs):
        # Default DecisionTreeClassifier params
        self.C=1.0
        self.kernel="rbf"
        self.degree=3
        self.gamma="auto"
        self.cache_size=200
        self.class_weight=None
        self.random_state = 123
        self.feats_to_ignore = []

        if "options" in kwargs:
            options = kwargs["options"]
            if "C" in options:
                self.C = options.pop("C")
            if "kernel" in options:
                self.kernel = options.pop("kernel")
            if "degree" in options:
                self.degree = options.pop("degree")
            if "gamma" in options:
                self.gamma = options.pop("gamma")
            if "random_state" in options:
                self.random_state = options.pop("random_state")
            if "cache_size" in options:
                self.cache_size = options.pop("cache_size")
            if "class_weight" in options:
                self.class_weight = options.pop("class_weight")
            if "feats_to_ignore" in options:
                self.feats_to_ignore = options.pop("feats_to_ignore")

        super(ModelFactory, self).__init__(*args, **kwargs)
        self.verbose_factory_name = "DecisionTree"

        # Maps each header to the values for it's column.
        self.col_vals = {header:{row[i] for row in self.all_data} for i, header in enumerate(self.headers)}
        self.response_index = self.headers.index(self.response_header)
        self.num_outcomes = len(self.col_vals[self.response_header])

        # Mark any categorical features for column expansion.
        # Categorical features are transfered into -1, 1 binary features for each value.
        self.columns_to_expand = []
        for i, header in enumerate(self.headers):
            categorical = any(type(val)==str for val in self.col_vals[header])
            if i == self.response_index or categorical:
                self.columns_to_expand.append(header)

        # All non-categorical features will be standardized (mean=0, variance=1).
        self.standardizers = {}
        for i, header in enumerate(self.headers):
            if any (type(row[i])==str for row in self.all_data):
                mean = None
                std_dev = None
            else:
                mean = sum(row[i] for row in self.all_data)/float(len(self.all_data))
                std_dev = np.sqrt(sum(((row[i] - mean) ** 2) for row in self.all_data)/float(len(self.all_data)))
            self.standardizers[header] = {'mean':mean, 'std_dev':std_dev}

        # If the response column is shifted by expanding categorical features,
        # update the response index.
        response_col_shift = 0
        for header in self.headers[:self.response_index]:
            if header in self.columns_to_expand:
                response_col_shift += len(self.col_vals[header]) - 1
        self.adjusted_response_index = self.response_index + response_col_shift

        # Map each outcome to an index
        self.outcome_trans_dict = {val:i for i, val in enumerate(self.col_vals[self.response_header])}

    def build(self, train_set):
        # prepare train data for learning
        expanded_and_stdized_train_set, self.expanded_headers = expand_and_standardize_dataset(self.response_index, self.response_header, train_set, self.col_vals, self.headers, self.standardizers, self.feats_to_ignore, self.columns_to_expand, self.outcome_trans_dict)
        train_matrix, train_outcomes = list_to_tf_input(expanded_and_stdized_train_set, self.adjusted_response_index, self.num_outcomes)
        train_size, num_features = train_matrix.shape

        clf = svm.SVC(
                kernel=self.kernel,\
                degree=self.degree,\
                gamma=self.gamma,\
                cache_size=self.cache_size,\
                class_weight=self.class_weight,\
                random_state=self.random_state
                )

        clf.fit(train_matrix, train_outcomes)
        model_name = "SVM_Model"

        return(ModelVisitor(clf, model_name, self.response_header, self.response_index, self.adjusted_response_index, self.num_outcomes, self.outcome_trans_dict, self.headers, self.expanded_headers, self.standardizers, self.col_vals, self.feats_to_ignore, self.columns_to_expand))

class ModelVisitor(AbstractModelVisitor):

  def __init__(self, clf, model_name, response_header, response_index, adjusted_response_index, num_outcomes, outcome_trans_dict, headers, expanded_headers, standardizers, train_col_vals, feats_to_ignore, columns_to_expand):
    super(ModelVisitor,self).__init__(model_name)
    self.response_index = response_index
    self.response_header = response_header
    self.adjusted_response_index = adjusted_response_index
    self.num_outcomes = num_outcomes
    self.outcome_trans_dict = outcome_trans_dict
    self.headers = headers
    self.expanded_headers = expanded_headers
    self.train_expanded_headers = expanded_headers
    self.train_col_vals = train_col_vals
    self.standardizers = standardizers
    self.feats_to_ignore = feats_to_ignore
    self.columns_to_expand = columns_to_expand
    self.clf = clf

  def test(self, test_set, test_name=""):
    expanded_and_stdized_test_set, self.test_set_expanded_headers = expand_and_standardize_dataset(self.response_index, self.response_header, test_set, self.train_col_vals, self.headers, self.standardizers, self.feats_to_ignore, self.columns_to_expand, self.outcome_trans_dict)
    if self.test_set_expanded_headers != self.train_expanded_headers:
      raise ValueError('Feature dimensions do not align!')
    else:
      test_matrix, test_labels = list_to_tf_input(expanded_and_stdized_test_set, self.adjusted_response_index, self.num_outcomes)

    predictions = self.clf.predict(test_matrix)
    predictions, test_labels = predictions.tolist(), test_labels.tolist() # cast from numpy array to python list
    #test_labels = [(corr[0], corr[1]) for corr in test_labels]
    #predictions = [row.index(1) for row in predictions]
    predictions_dict = {i:key for key,i in list(self.outcome_trans_dict.items())}
    predictions = [predictions_dict[pred] for pred in predictions]
    return list(zip([row[self.response_index] for row in test_set], predictions))

def list_to_tf_input(data, response_index, num_outcomes):
  """
  Separates the outcome feature from the data.
  """
  matrix = np.matrix([row[:response_index] + row[response_index+1:] for row in data])
  outcomes = np.asarray([row[response_index] for row in data], dtype=np.uint8)

  return matrix, outcomes

def expand_and_standardize_dataset(response_index, response_header, data_set, col_vals, headers, standardizers, feats_to_ignore, columns_to_expand, outcome_trans_dict):
  """
  Standardizes continuous features and expands categorical features.
  """
  # expand and standardize
  modified_set = []
  for row_index, row in enumerate(data_set):
    new_row = []
    for col_index, val in enumerate(row):
      header = headers[col_index]

      # Outcome feature -> index outcome
      if col_index == response_index:
        new_outcome = outcome_trans_dict[val]
        new_row.append(new_outcome)

      # Ignored feature -> pass
      elif header in feats_to_ignore:
        pass

      # Categorical feature -> create new binary column for each possible value of the column
      elif header in columns_to_expand:
        for poss_val in col_vals[header]:
          if val == poss_val:
            new_cat_val = 1.0
          else:
            new_cat_val = -1.0
          new_row.append(new_cat_val)

      # Continuous feature -> standardize value with respect to its column
      else:
        new_cont_val = float((val - standardizers[header]['mean']) / standardizers[header]['std_dev'])
        new_row.append(new_cont_val)

    modified_set.append(new_row)

  # update headers to reflect column expansion
  expanded_headers = []
  for header in headers:
    if header in feats_to_ignore:
      pass
    elif (header in columns_to_expand) and (header is not response_header):
      for poss_val in col_vals[header]:
        new_header = '{}_{}'.format(header,poss_val)
        expanded_headers.append(new_header)
    else:
      expanded_headers.append(header)

  return modified_set, expanded_headers

def test():
  test_expand_and_standardize_dataset()
  test_unseen_categorical_feature()
  test_categorical_model()
  test_categorical_response()
  test_list_to_tf_input()
  test_basic_model()
  test_simple_data()

def test_expand_and_standardize_dataset():
  """
  Tests if the function correctly converts categorical
  features into onehot features, correctly standardizes
  continuous features, and correctly converts the response
  feature to indexes.
  """
  headers = ['cont_feature', 'cat_feature', 'ignored', 'response']
  data = [[75., 'A', 'ingore1', 'X'], [25, 'B', 'ignore2', 'Y'], [75, 'C', 'ignore3', 'Z'], [25., 'A', 'ignore1', 'Y']]
  resp_index = 3
  resp_header = 'response'
  feats_to_ignore = 'ignored'
  col_vals = {'cont_feature':[75., 25.], 'cat_feature':['A','B','C'], 'ignored':['ignore1', 'ignore2', 'ignore3'], 'response':['X', 'Y', 'Z']}
  columns_to_expand = ['cat_feature', 'ignored', 'response']
  outcome_trans_dict = {'X':0, 'Y':1, 'Z':2}
  standardizers = {'cont_feature':{'mean':50, 'std_dev':25}}
  new_data, new_headers = expand_and_standardize_dataset(resp_index, resp_header, data, col_vals, headers, standardizers, feats_to_ignore, columns_to_expand, outcome_trans_dict)
  correct_data = [[1.0, 1.0, -1.0, -1.0, 0], [-1.0, -1.0, 1.0, -1.0, 1], [1.0, -1.0, -1.0, 1.0, 2], [-1.0, 1.0, -1.0, -1.0, 1]]
  correct_headers = ['cont_feature', 'cat_feature_A', 'cat_feature_B', 'cat_feature_C', 'response']
  print('expanding and standardizing data correctly? --', new_data==correct_data)
  print('expanding headers correctly? --', new_headers==correct_headers)

def test_unseen_categorical_feature():
  """
  Tests if the model can correctly handle unfamiliar values
  for categorical features. The model should not create new onehot
  features for values in the test set and should instead treat
  these values as 'not any known value for that column'.
  """
  headers = ["predictor 1", "predictor 2", "predictor3", "response"]
  response = "response"
  random_features = ['H', 5.5, 'I']
  train_set = [[random_features[i%3],'B',1,'C'] for i in range(1,50)] + [['A','A',-1,'D'] for i in range(1,50)]
  train_set_copy = copy.copy(train_set)
  new_features = ['J',9]
  test_set = [[new_features[i%2],'B',1,'C'] for i in range(1,49)] + [['A','A',-1,'D'] for i in range(1,50)]
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")
  model = factory.build(train_set)
  print("factory builds ModelVisitor? -- ", isinstance(model, ModelVisitor))

  predictions = model.test(test_set)
  print("handling unseen categorical values correctly? -- ", all([pred[0] == pred[1] for pred in predictions]))

def test_list_to_tf_input():
  data = [[0,0],[1,3],[0,2]]
  tf_matrix, tf_output = list_to_tf_input(data, 1, 3)
  correct_matrix = [[0],[1],[0]]
  correct_output = [0,3,2]
  print("list_to_tf_input matrix correct? --",np.array_equal(tf_matrix, correct_matrix))
  print("list_to_tf_input onehot correct? --",np.array_equal(tf_output, correct_output))

def test_basic_model():
  headers = ["predictor 1", "predictor 2", "response"]
  response = "response"
  train_set = [[i,0,0] for i in range(1,50)] + [[0,i,1] for i in range(1,50)]
  test_set = [[i,0,0] for i in range(1,50)] + [[0,i,1] for i in range(1,50)]
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")

  model = factory.build(train_set)

  predictions = model.test(test_set)
  intended_predictions = [pred[0] for pred in predictions]
  actual_predictions = [pred[1] for pred in predictions]
  print("predicting numeric categories correctly? -- ", actual_predictions == intended_predictions)

def test_categorical_response():
  headers = ["predictor 1", "predictor 2", "response"]
  response = "response"
  train_set = [[i,0,"A"] for i in range(1,50)] + [[0,i,"B"] for i in range(1,50)]
  test_set = [[i,0,"A"] for i in range(1,50)] + [[0,i,"C"] for i in range(1,50)]
  train_set_copy = copy.copy(train_set)
  test_set_copy = copy.copy(test_set)
  all_data = train_set + test_set

  factory = ModelFactory(all_data, headers, response, name_prefix="test")
  model = factory.build(train_set)

  predictions = model.test(test_set)
  resp_index = headers.index(response)
  intended_predictions = [(test_row[resp_index], train_row[resp_index]) for train_row, test_row in zip(train_set_copy,test_set_copy)]
  print("predicting string-categories correctly? -- ", predictions == intended_predictions)

def test_categorical_model():
  """
  Tests whether the model can correctly learn on the train set.
  The model should have poor accuracy on the test set
  because the outcomes do not match those of the train set.
  The model should incorrectly predict response 'B' for predictor 'B'
  when response 'C' is the correct response.
  """
  headers = ["predictor", "response"]
  response = "response"
  train_set = [["A","A"] for i in range(1,50)] + [["B","B"] for i in range(1,50)]
  train_set_copy = copy.copy(train_set)
  test_set = [["A","A"] for i in range(1,50)] + [["B","C"] for i in range(1,50)]
  all_data = train_set + test_set
  factory = ModelFactory(all_data, headers, response, name_prefix="test", options={})

  model = factory.build(train_set)

  predictions = model.test(test_set)
  resp_index = headers.index(response)
  intended_predictions = [(test_row[resp_index], train_row[resp_index]) for train_row, test_row in zip(train_set_copy,test_set)]
  print("predicting string-categories correctly? -- ", predictions == intended_predictions)

def test_simple_data():
    headers = ["predictor_1", "predictor_2", "response"]
    response = "response"

    train_set = [["A","B","C"] for i in range(1,50)] + [["B","A","D"] for i in range(1,50)] + [["B","B","E"] for i in range(1,50)] + [["A", "A", "E"] for i in range(1,50)] + [["C", "D", "F"] for i in range(1,10)]
    train_set_copy = copy.copy(train_set)
    test_set = [["A","B","C"] for i in range(1,50)] + [["B","A","D"] for i in range(1,50)] + [["B","B","E"] for i in range(1,50)] + [["A", "A", "E"] for i in range(1,50)] + [["A","X","Y"] for i in range(1,10)]
    all_data = train_set + test_set
    factory = ModelFactory(all_data, headers, response, name_prefix="test", options={})

    model = factory.build(train_set)
    predictions = model.test(test_set)
    resp_index = headers.index(response)

    intended_predictions = [(test_row[resp_index], train_row[resp_index]) for train_row, test_row in zip(train_set_copy,test_set)]
    last_few_preds = all(pred[0] == pred[1] for pred in predictions[-9:])
    all_but_last_few = all(pred[0] == pred[1] for pred in predictions[:-9])
    print("handling simple data correctly? -- ", (all_but_last_few and not last_few_preds) == True)

if __name__=='__main__':
  test()


import numpy

def FreedmanDiaconisBinSize(feature_values):
  """
  The bin size in FD-binning is given by size = 2 * IQR(x) * n^(-1/3)
  More Info: https://en.wikipedia.org/wiki/Freedman%E2%80%93Diaconis_rule

  If the BinSize ends up being 0 (in the case that all values are the same),
  return a BinSize of 1.
  """

  q75, q25 = numpy.percentile(feature_values, [75, 25])
  IQR = q75 - q25

  return 2.0 * IQR * len(feature_values) ** (-1.0/3.0)

def test():
  values = range(0,100)
  bin_size = FreedmanDiaconisBinSize(values)
  correct_bin_size = 21
  bin_size = round(bin_size)
  print "FreedmanDiaconisBinSize -- correct size of bins? ", bin_size == correct_bin_size

if __name__=="__main__":
  test()


from __future__ import print_function, unicode_literals

import os

import re
import requests
import sys

from colorclass import Color
from packaging import version
from packaging.utils import canonicalize_name
from requests import HTTPError

try:
    from pip.locations import site_config_files
except ImportError:
    from pip._internal.locations import site_config_files

try:
    from configparser import ConfigParser, NoOptionError, NoSectionError
except ImportError:   # pragma: nocover
    from ConfigParser import ConfigParser, NoOptionError, NoSectionError

try:
    from urllib.parse import urljoin
except ImportError:  # pragma: nocover
    from urlparse import urljoin

class PackagesStatusDetector(object):
    packages = []
    packages_status_map = {}
    PYPI_API_URL = None
    PYPI_API_TYPE = None
    pip_config_locations = [
        '~/.pip/pip.conf',
        '~/.pip/pip.ini',
        '~/.config/pip/pip.conf',
        '~/.config/pip/pip.ini',
    ]

    _prerelease = False

    def __init__(self, packages, use_default_index=False):
        self.packages = packages
        self.packages_status_map = {}
        self.PYPI_API_URL = 'https://pypi.python.org/pypi/{package}/json'
        self.PYPI_API_TYPE = 'pypi_json'

        if not use_default_index:
            self._update_index_url_from_configs()

        self._prerelease = False

    def _update_index_url_from_configs(self):
        """ Checks for alternative index-url in pip.conf """

        if 'VIRTUAL_ENV' in os.environ:
            self.pip_config_locations.append(os.path.join(os.environ['VIRTUAL_ENV'], 'pip.conf'))
            self.pip_config_locations.append(os.path.join(os.environ['VIRTUAL_ENV'], 'pip.ini'))

        if site_config_files:
            self.pip_config_locations.extend(site_config_files)

        index_url = None
        custom_config = None

        if 'PIP_INDEX_URL' in os.environ and os.environ['PIP_INDEX_URL']:
            # environ variable takes priority
            index_url = os.environ['PIP_INDEX_URL']
            custom_config = 'PIP_INDEX_URL environment variable'
        else:
            for pip_config_filename in self.pip_config_locations:
                if pip_config_filename.startswith('~'):
                    pip_config_filename = os.path.expanduser(pip_config_filename)

                if os.path.isfile(pip_config_filename):
                    config = ConfigParser()
                    config.read([pip_config_filename])
                    try:
                        index_url = config.get('global', 'index-url')
                        custom_config = pip_config_filename
                        break  # stop on first detected, because config locations have a priority
                    except (NoOptionError, NoSectionError):  # pragma: nocover
                        pass

        if index_url:
            self.PYPI_API_URL = self._prepare_api_url(index_url)
            print(Color('Setting API url to {{autoyellow}}{}{{/autoyellow}} as found in {{autoyellow}}{}{{/autoyellow}}'
                        '. Use --default-index-url to use pypi default index'.format(self.PYPI_API_URL, custom_config)))

    def _prepare_api_url(self, index_url):  # pragma: nocover
        if not index_url.endswith('/'):
            index_url += '/'

        if index_url.endswith('/simple/'):
            self.PYPI_API_TYPE = 'simple_html'
            return urljoin(index_url, '{package}')

        if index_url.endswith('/+simple/'):
            self.PYPI_API_TYPE = 'simple_html'
            return urljoin(index_url, '{package}')

        if '/pypi/' in index_url:
            base_url = index_url.split('/pypi/')[0]
            return urljoin(base_url, '/pypi/{package}/json')

        return urljoin(index_url, '/pypi/{package}/json')

    def detect_available_upgrades(self, options):
        self._prerelease = options.get('--prerelease', False)
        explicit_packages_lower = None
        if options['-p'] and options['-p'] != ['all']:
            explicit_packages_lower = [pack_name.lower() for pack_name in options['-p']]

        for i, package in enumerate(self.packages):
            try:
                package_name, pinned_version = self._expand_package(package)
                if not package_name or not pinned_version:  # pragma: nocover
                    # todo: treat <= or >= instead of ==
                    continue

                if explicit_packages_lower and package_name.lower() not in explicit_packages_lower:
                    # skip if explicit and not chosen
                    continue

                current_version = version.parse(pinned_version)

                if pinned_version and isinstance(current_version, version.Version):  # version parsing is correct
                    package_status, reason = self._fetch_index_package_info(package_name, current_version)
                    if not package_status:  # pragma: nocover
                        print(package, reason)
                        continue

                    print('{}/{}: {} ... '.format(i + 1, len(self.packages), package_name), end='')
                    sys.stdout.flush()

                    # compare versions
                    if current_version < package_status['latest_version']:
                        print('upgrade available: {} ==> {} (uploaded on {})'.format(current_version,
                                                                                package_status['latest_version'],
                                                                                package_status['upload_time']))
                    else:
                        print('up to date: {}'.format(current_version))
                    sys.stdout.flush()

                    self.packages_status_map[package_name] = package_status
            except Exception as e:  # noqa  # pragma: nocover
                print('Error while parsing package {} (skipping). \nException: '.format(package), e)

        return self.packages_status_map

    def _fetch_index_package_info(self, package_name, current_version):
        """
        :type package_name: str
        :type current_version: version.Version
        """

        try:
            package_canonical_name = package_name
            if self.PYPI_API_TYPE == 'simple_html':
                package_canonical_name = canonicalize_name(package_name)
            response = requests.get(self.PYPI_API_URL.format(package=package_canonical_name), timeout=15)
        except HTTPError as e:  # pragma: nocover
            return False, e.message

        if not response.ok:  # pragma: nocover
            return False, 'API error: {}'.format(response.reason)

        if self.PYPI_API_TYPE == 'pypi_json':
            return self._parse_pypi_json_package_info(package_name, current_version, response)
        elif self.PYPI_API_TYPE == 'simple_html':
            return self._parse_simple_html_package_info(package_name, current_version, response)
        else:  # pragma: nocover
            raise NotImplementedError('This type of PYPI_API_TYPE type is not supported')

    def _expand_package(self, package_line):
        if '==' in package_line:
            name, vers = package_line.split('==')

            if '[' in name and name.strip().endswith(']'):
                name = name.split('[')[0]

            return name, vers

        return None, None

    def _parse_pypi_json_package_info(self, package_name, current_version, response):
        """
        :type package_name: str
        :type current_version: version.Version
        :type response: requests.models.Response
        """

        data = response.json()
        all_versions = [version.parse(vers) for vers in data['releases'].keys()]
        filtered_versions = [vers for vers in all_versions if not vers.is_prerelease and not vers.is_postrelease]

        if not filtered_versions:  # pragma: nocover
            return False, 'error while parsing version'

        latest_version = max(filtered_versions)

        # even if user did not choose prerelease, if the package from requirements is pre/post release, use it
        if self._prerelease or current_version.is_postrelease or current_version.is_prerelease:
            prerelease_versions = [vers for vers in all_versions if vers.is_prerelease or vers.is_postrelease]
            if prerelease_versions:
                latest_version = max(prerelease_versions)
        try:
            try:
                latest_version_info = data['releases'][str(latest_version)][0]
            except KeyError:  # pragma: nocover
                # non-RFC versions, get the latest from pypi response
                latest_version = version.parse(data['info']['version'])
                latest_version_info = data['releases'][str(latest_version)][0]
        except Exception:  # pragma: nocover
            return False, 'error while parsing version'

        upload_time = latest_version_info['upload_time'].replace('T', ' ')

        return {
            'name': package_name,
            'current_version': current_version,
            'latest_version': latest_version,
            'upgrade_available': current_version < latest_version,
            'upload_time': upload_time
        }, 'success'

    def _parse_simple_html_package_info(self, package_name, current_version, response):
        """
        :type package_name: str
        :type current_version: version.Version
        :type response: requests.models.Response
        """
        pattern = r'<a.*>.*{name}-([A-z0-9\.-]*)(?:-py|\.tar).*<\/a>'.format(name=re.escape(package_name))
        versions_match = re.findall(pattern, response.content.decode('utf-8'), flags=re.IGNORECASE)

        all_versions = [version.parse(vers) for vers in versions_match]
        filtered_versions = [vers for vers in all_versions if not vers.is_prerelease and not vers.is_postrelease]

        if not filtered_versions:  # pragma: nocover
            return False, 'error while parsing version'

        latest_version = max(filtered_versions)

        # even if user did not choose prerelease, if the package from requirements is pre/post release, use it
        if self._prerelease or current_version.is_postrelease or current_version.is_prerelease:
            prerelease_versions = [vers for vers in all_versions if vers.is_prerelease or vers.is_postrelease]
            if prerelease_versions:
                latest_version = max(prerelease_versions)

        return {
           'name': package_name,
           'current_version': current_version,
           'latest_version': latest_version,
           'upgrade_available': current_version < latest_version,
           'upload_time': '-'
        }, 'success'


"""
pip-upgrade

Usage:
  pip-upgrade [<requirements_file>] ... [--prerelease] [-p=<package>...] [--dry-run] [--skip-virtualenv-check] [--skip-package-installation] [--use-default-index]

Arguments:
    requirements_file             The requirement FILE, or WILDCARD PATH to multiple files.
    --prerelease                  Include prerelease versions for upgrade, when querying pypi repositories.
    -p <package>                  Pre-choose which packages tp upgrade. Skips any prompt.
    --dry-run                     Simulates the upgrade, but does not execute the actual upgrade.
    --skip-package-installation   Only upgrade the version in requirements files, don't install the new package.
    --skip-virtualenv-check       Disable virtualenv check. Allows installing the new packages outside the virtualenv.
    --use-default-index           Skip searching for custom index-url in pip configuration file(s).

Examples:
  pip-upgrade             # auto discovers requirements file
  pip-upgrade requirements.txt
  pip-upgrade requirements/dev.txt requirements/production.txt
  pip-upgrade requirements.txt -p django -p celery
  pip-upgrade requirements.txt -p all
  pip-upgrade requirements.txt --dry-run  # run everything as a simulation (don't do the actual upgrade)

Help:
  Interactively upgrade packages from requirements file, and also update the pinned version from requirements file(s).
  If no requirements are given, the command attempts to detect the requirements file(s) in the current directory.

  https://github.com/simion/pip-upgrader
"""  # noqa: E501
from __future__ import print_function, unicode_literals
from colorclass import Windows, Color
from docopt import docopt

from pip_upgrader import __version__ as VERSION
from pip_upgrader.packages_detector import PackagesDetector
from pip_upgrader.packages_interactive_selector import PackageInteractiveSelector
from pip_upgrader.packages_status_detector import PackagesStatusDetector
from pip_upgrader.packages_upgrader import PackagesUpgrader
from pip_upgrader.requirements_detector import RequirementsDetector
from pip_upgrader.virtualenv_checker import check_for_virtualenv

def get_options():
    return docopt(__doc__, version=VERSION)

def main():
    """ Main CLI entrypoint. """
    options = get_options()
    Windows.enable(auto_colors=True, reset_atexit=True)

    try:
        # maybe check if virtualenv is not activated
        check_for_virtualenv(options)

        # 1. detect requirements files
        filenames = RequirementsDetector(options.get('<requirements_file>')).get_filenames()
        if filenames:
            print(Color('{{autoyellow}}Found valid requirements file(s):{{/autoyellow}} '
                        '{{autocyan}}\n{}{{/autocyan}}'.format('\n'.join(filenames))))
        else:  # pragma: nocover
            print(Color('{autoyellow}No requirements files found in current directory. CD into your project '
                        'or manually specify requirements files as arguments.{/autoyellow}'))
            return
        # 2. detect all packages inside requirements
        packages = PackagesDetector(filenames).get_packages()

        # 3. query pypi API, see which package has a newer version vs the one in requirements (or current env)
        packages_status_map = PackagesStatusDetector(
            packages, options.get('--use-default-index')).detect_available_upgrades(options)

        # 4. [optionally], show interactive screen when user can choose which packages to upgrade
        selected_packages = PackageInteractiveSelector(packages_status_map, options).get_packages()

        # 5. having the list of packages, do the actual upgrade and replace the version inside all filenames
        upgraded_packages = PackagesUpgrader(selected_packages, filenames, options).do_upgrade()

        print(Color('{{autogreen}}Successfully upgraded (and updated requirements) for the following packages: '
                    '{}{{/autogreen}}'.format(','.join([package['name'] for package in upgraded_packages]))))
        if options['--dry-run']:
            print(Color('{automagenta}Actually, no, because this was a simulation using --dry-run{/automagenta}'))

    except KeyboardInterrupt:  # pragma: nocover
        print(Color('\n{autored}Upgrade interrupted.{/autored}'))


from __future__ import print_function, unicode_literals

import subprocess
from subprocess import CalledProcessError

import re
from colorclass import Color

class PackagesUpgrader(object):

    selected_packages = None
    requirements_files = None
    upgraded_packages = None
    dry_run = False

    def __init__(self, selected_packages, requirements_files, options):
        self.selected_packages = selected_packages
        self.requirements_files = requirements_files
        self.upgraded_packages = []
        self.dry_run = options['--dry-run']
        self.skip_package_installation = options.get('--skip-package-installation', False)

    def do_upgrade(self):
        for package in self.selected_packages:
            self._update_package(package)

        return self.upgraded_packages

    def _update_package(self, package):
        """ Update (install) the package in current environment, and if success, also replace version in file """
        try:
            if not self.dry_run and not self.skip_package_installation:  # pragma: nocover
                subprocess.check_call(['pip', 'install', '{}=={}'.format(package['name'], package['latest_version'])])
            else:
                print('[Dry Run]: skipping package installation:', package['name'])
            # update only if installation success
            self._update_requirements_package(package)

        except CalledProcessError:  # pragma: nocover
            print(Color('{{autored}}Failed to install package "{}"{{/autored}}'.format(package['name'])))

    def _update_requirements_package(self, package):
        for filename in set(self.requirements_files):
            lines = []

            # read current lines
            with open(filename, 'r') as frh:
                for line in frh:
                    lines.append(line)

            try:
                # write updates lines
                with open(filename, 'w') as fwh:
                    for line in lines:
                        line = self._maybe_update_line_package(line, package)
                        fwh.write(line)
            except Exception as e:  # pragma: nocover
                # at exception, restore old file
                with open(filename, 'w') as fwh:
                    for line in lines:
                        fwh.write(line)
                raise e

    def _maybe_update_line_package(self, line, package):
        original_line = line
        pattern = r'\b{package}(?:\[\w*\])?=={old_version}\b'.format(
            package=re.escape(package['name']),
            old_version=re.escape(str(package['current_version'])))

        if re.search(pattern, line, flags=re.IGNORECASE):
            line = line.replace(
                '=={}'.format(package['current_version']),
                '=={}'.format(package['latest_version'])
            )

        if line != original_line:
            self.upgraded_packages.append(package)

            if self.dry_run:  # pragma: nocover
                print('[Dry Run]: skipping requirements replacement:',
                      original_line.replace('\n', ''), ' / ',
                      line.replace('\n', ''))
                return original_line
        return line


from __future__ import print_function, unicode_literals

import mimetypes
import os

class RequirementsDetector(object):
    """ Takes raw requirements argument, and detects / discovers all the requirements files. """

    filenames = []

    def __init__(self, requirements_arg):
        self.filenames = []

        if not requirements_arg:
            self.autodetect_files()
        else:
            self.detect_files(requirements_arg)

    def get_filenames(self):
        """ Returns a list of all filenames detected as proper requirements files. """
        return self.filenames

    def detect_files(self, requirements_arg):
        for argument in requirements_arg:
            if self._is_valid_requirements_file(argument):
                self.filenames.append(argument)
            else:  # pragma: nocover
                print('Invalid requirements file: {}'.format(argument))
        self._check_inclusions_recursively()

    def autodetect_files(self):
        """ Attempt to detect requirements files in the current working directory """
        if self._is_valid_requirements_file('requirements.txt'):
            self.filenames.append('requirements.txt')

        if self._is_valid_requirements_file('requirements.pip'):  # pragma: nocover
            self.filenames.append('requirements.pip')

        if os.path.isdir('requirements'):
            for filename in os.listdir('requirements'):
                file_path = os.path.join('requirements', filename)
                if self._is_valid_requirements_file(file_path):
                    self.filenames.append(file_path)
        self._check_inclusions_recursively()

    @staticmethod
    def _is_valid_requirements_file(filename):
        extension_ok = filename.endswith('txt') or filename.endswith('pip')
        return extension_ok and os.path.isfile(filename) and mimetypes.guess_type(filename)[0] in ['text/plain', None]

    def _check_inclusions_recursively(self):
        for filename in self.filenames:
            self._detect_inclusion(filename)

    def _detect_inclusion(self, filename):
        with open(filename) as fh:
            for line in fh:
                if line.strip().startswith('-r '):
                    included_filename = line.split('-r ')[1].strip()
                    included_filename = os.path.join(os.path.dirname(filename), included_filename)
                    if self._is_valid_requirements_file(included_filename) and included_filename not in self.filenames:
                        self.filenames.append(included_filename)
                        # recursively, check if the included file contains other inclusions
                        self._detect_inclusion(included_filename)