# This code should be run with the library from
# https://github.com/mbilalzafar/fair-classification/

import os,sys
import numpy as np
sys.path.insert(0, '../../fair_classification/') # the code for fair classification is in this directory
import utils as ut
import loss_funcs as lf # loss funcs that can be optimized subject to various constraints
import csv
import random
import scipy
import scipy.stats
import math
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from sklearn import metrics
from collections import defaultdict

def sigmoid(x):
  if x > 50:
    x = 50
  if x < -50:
    x = -50
  return 1.0 / (1 + math.exp(-x))

def readData(fname):
  f = open(fname, 'r')
  X_raw = [] # Raw features (with sensitive variable, without label). For debugging only
  X = [] # Reduced set of features to be used for prediction
  X_with_sens = [] # X with the sensitive feature
  x_control = {"s1": []} # Sensitive feature
  y = [] # Label

  # Functions to read the various datasets...
  if "adult.data" in fname: # UCI-Income
    reader = csv.reader(f)
    header = ["age", "workclass", "fnlwgt", "education", "education-num", "marital-status", "occupation", "relationship", "race", "sex", "capital-gain", "capital-loss", "hours-per-week", "native-country", "income"]
    fieldIds = dict(zip(header,range(len(header))))

    rows = []
    for row in reader:
      rows.append([x.strip() for x in row])
    random.shuffle(rows)

    for row in rows:
      d = dict(zip(header,row))
      try:
        x1 = [1, float(d['age'])]
      except Exception as e:
        continue # Empty last line
      x1 += [float(d['workclass'] == c) for c in ["Private", "Self-emp-not-inc", "Self-emp-inc", "Federal-gov", "Local-gov", "State-gov", "Without-pay", "Never-worked"]]
      x1.append(float(d['fnlwgt']))
      x1.append(float(d['education-num']))
      x1 += [float(d['marital-status'] == c) for c in ["Married-civ-spouse", "Divorced", "Never-married", "Separated", "Widowed", "Married-spouse-absent", "Married-AF-spouse"]]
      x1 += [float(d['occupation'] == c) for c in ["Tech-support", "Craft-repair", "Other-service", "Sales", "Exec-managerial", "Prof-specialty", "Handlers-cleaners", "Machine-op-inspct", "Adm-clerical", "Farming-fishing", "Transport-moving", "Priv-house-serv", "Protective-serv", "Armed-Forces"]]
      x1 += [float(d['relationship'] == c) for c in ["Wife", "Own-child", "Husband", "Not-in-family", "Other-relative", "Unmarried"]]
      x1 += [float(d['race'] == c) for c in ["White", "Asian-Pac-Islander", "Amer-Indian-Eskimo", "Other", "Black"]]
      x1.append(float(d['capital-gain']))
      x1.append(float(d['capital-loss']))
      x1.append(float(d['hours-per-week']))
      x1 += [float(d['native-country'] == c) for c in ['United-States', 'Cambodia', 'England', 'Puerto-Rico', 'Canada', 'Germany', 'Outlying-US(Guam-USVI-etc)', 'India', 'Japan', 'Greece', 'South', 'China', 'Cuba', 'Iran', 'Honduras', 'Philippines', 'Italy', 'Poland', 'Jamaica', 'Vietnam', 'Mexico', 'Portugal', 'Ireland', 'France', 'Dominican-Republic', 'Laos', 'Ecuador', 'Taiwan', 'Haiti', 'Columbia', 'Hungary', 'Guatemala', 'Nicaragua', 'Scotland', 'Thailand', 'Yugoslavia', 'El-Salvador', 'Trinadad&Tobago', 'Peru', 'Hong', 'Holand-Netherlands']]

      x1 += [1 * d['sex'] == 'Female']
      sensitiveId = len(x1) - 1

      x_with_sens = x1
      X_raw.append(x1)
      X.append(x1[:-1])
      X_with_sens.append(x_with_sens)
      x_control["s1"].append(1 - x1[sensitiveId])
      y.append(int(d['income'] == '>50K')*2 - 1)

  if "bank-full.csv" in fname: # UCI-Marketing
    reader = csv.reader(f, delimiter=';')
    header = [x.replace('"', '') for x in reader.next()]
    fieldIds = dict(zip(header,range(len(header))))

    rows = []
    for row in reader:
      row = [x.replace('"', '') for x in row]
      rows.append(row)
    random.shuffle(rows)

    for row in rows:
      d = dict(zip(header,row))
      x1 = [1]
      x1 += [float(d['age'])]
      x1 += [float(d['job'] == c) for c in ['admin.','blue-collar','entrepreneur','housemaid','management','retired','self-employed','services','student','technician','unemployed','unknown']]
      x1 += [float(d['education'] == c) for c in ['basic.4y','basic.6y','basic.9y','high.school','illiterate','professional.course','university.degree','unknown']]
      x1.append(1 * (d['default'] == 'yes'))
      x1.append(1 * (d['housing'] == 'yes'))
      x1.append(1 * (d['loan'] == 'yes'))
      x1.append(1 * (d['contact'] == 'cellular'))
      x1 += [float(d['month'] == c) for c in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul', 'aug', 'sep', 'oct', 'nov', 'dec']]
      x1 += [float(d['day'] == c) for c in ['mon', 'tue', 'wed', 'thu', 'fri']]
      x1.append(float(d['campaign']))
      x1.append(float(d['pdays']))
      x1.append(float(d['previous']))
      x1.append(1 * (d['poutcome'] == 'success'))

      x1 += [1 * d['marital'] == 'married']
      sensitiveId = len(x1) - 1

      x_with_sens = x1
      X_raw.append(x1)
      X.append(x1[:-1])
      X_with_sens.append(x_with_sens)
      x_control["s1"].append(1 - x1[sensitiveId])
      y.append(int(d['y'] == 'yes')*2 - 1)

  if "default.csv" in fname: # UCI-Credit
    reader = csv.reader(f, delimiter=',')
    header = reader.next()
    fieldIds = dict(zip(header,range(len(header))))

    rows = []
    for row in reader:
      rows.append(row)
    random.shuffle(rows)

    for row in rows:
      d = dict(zip(header,row))
      x1 = [float(d['LIMIT_BAL'])]
      x1 += [float(d['EDUCATION'] == c) for c in ['1', '2', '3', '4']]
      x1 += [float(d['MARRIAGE'] == c) for c in ['1', '2', '3']]
      x1.append(float(d['AGE']))
      for field in ['PAY_0', 'PAY_2', 'PAY_3', 'PAY_4', 'PAY_5', 'PAY_6']:
        x1 += [1 * (d[field] == '-2'), 1 * (d[field] == '-1'), max(0, float(d[field]))]
      for field in ['BILL_AMT1','BILL_AMT2','BILL_AMT3','BILL_AMT4','BILL_AMT5','BILL_AMT6','PAY_AMT1','PAY_AMT2','PAY_AMT3','PAY_AMT4','PAY_AMT5','PAY_AMT6']:
        x1.append(float(d[field]))


      x1 += [1 * (d['SEX'] == '2')]
      sensitiveId = len(x1) - 1

      x_with_sens = x1
      X_raw.append(x1)
      X.append(x1[:-1])
      X_with_sens.append(x_with_sens)
      x_control["s1"].append(1 - x1[sensitiveId])
      y.append(int(d['default payment next month'] == '1')*2 - 1)

  if "ibm_eval.csv" in fname: # IBM-Employee
    reader = csv.reader(f, delimiter=',')
    header = reader.next()
    fieldIds = dict(zip(header,range(len(header))))

    rows = []
    for row in reader:
      rows.append(row)
    random.shuffle(rows)

    for row in rows:
      if len(row) == 0:
        continue # Empty line?
      d = dict(zip(header,row))
      x1 = [1]
      x1.append(float(d['Age']))
      x1 += [float(d['BusinessTravel'] == c) for c in ['Travel_Rarely', 'Travel_Frequently', 'Non-Travel']]
      x1.append(float(d['DailyRate']))
      x1 += [float(d['Department'] == c) for c in ['Sales', 'Research & Development', 'Human Resources']]
      x1.append(float(d['DistanceFromHome']))
      x1 += [float(d['Education'] == c) for c in ['1', '2', '3', '4', '5']]
      x1 += [float(d['EnvironmentSatisfaction'] == c) for c in ['1', '2', '3', '4']]
      x1 += [float(d['JobInvolvement'] == c) for c in ['1', '2', '3', '4']]
      x1 += [float(d['JobSatisfaction'] == c) for c in ['1', '2', '3', '4']]
      x1 += [float(d['RelationshipSatisfaction'] == c) for c in ['1', '2', '3', '4']]
      x1 += [float(d['WorkLifeBalance'] == c) for c in ['1', '2', '3', '4']]
      x1.append(float(d['HourlyRate']))
      x1.append(float(d['JobLevel']))
      x1.append(float(d['MonthlyIncome']))
      x1.append(float(d['MonthlyRate']))
      x1.append(float(d['NumCompaniesWorked']))
      x1.append(1 * (d['Over18'] == 'Y'))
      x1.append(1 * (d['OverTime'] == 'Yes'))
      x1.append(float(d['PercentSalaryHike']))
      x1.append(float(d['StandardHours']))
      x1.append(float(d['StockOptionLevel']))
      x1.append(float(d['TotalWorkingYears']))
      x1.append(float(d['TrainingTimesLastYear']))
      x1.append(float(d['YearsAtCompany']))
      x1.append(float(d['YearsInCurrentRole']))
      x1.append(float(d['YearsSinceLastPromotion']))
      x1.append(float(d['YearsWithCurrManager']))
      x1.append(1 * (d['Gender'] == 'Female'))


      x1 += [1 * (d['MaritalStatus'] == 'Married')]
      sensitiveId = len(x1) - 1

      x_with_sens = x1
      X_raw.append(x1)
      X.append(x1[:-1])
      X_with_sens.append(x_with_sens)
      x_control["s1"].append(1 - x1[sensitiveId])
      y.append(int(d['Attrition'] == 'Yes')*2 - 1)

  if "ibm_churn.csv" in fname: # IBM-Customer
    reader = csv.reader(f, delimiter=',')
    header = reader.next()
    fieldIds = dict(zip(header,range(len(header))))

    rows = []
    for row in reader:
      rows.append(row)
    random.shuffle(rows)

    for row in rows:
      if len(row) == 0:
        continue # Empty line?
      d = dict(zip(header,row))
      x1 = [1]
      x1.append(1 * (d['SeniorCitizen'] == '1'))
      x1.append(1 * (d['Dependents'] == 'Yes'))
      x1.append(float(d['tenure']))
      x1.append(1 * (d['PhoneService'] == 'Yes'))
      x1 += [float(d['MultipleLines'] == c) for c in ['No', 'Yes', 'No phone service']]
      x1 += [float(d['InternetService'] == c) for c in ['DSL', 'Fiber optic', 'No']]
      for field in ['OnlineSecurity', 'OnlineBackup', 'DeviceProtection', 'TechSupport', 'StreamingTV', 'StreamingMovies']:
        x1 += [float(d[field] == c) for c in ['No', 'Yes', 'No internet service']]
      x1 += [float(d['Contract'] == c) for c in ['Month-to-month', 'One year', 'Two year']]
      x1.append(1 * (d['PaperlessBilling'] == 'Yes'))
      x1 += [float(d['PaymentMethod'] == c) for c in ['Electronic check', 'Mailed check', 'Bank transfer (automatic)', 'Credit card (automatic)']]
      x1.append(float(d['MonthlyCharges']))
      try:
        x1.append(float(d['TotalCharges']))
      except Exception as e:
        x1.append(0)
      x1.append(1 * (d['gender'] == 'Female'))

      x1 += [1 * (d['Partner'] == 'Yes')]
      sensitiveId = len(x1) - 1

      x_with_sens = x1
      X_raw.append(x1)
      X.append(x1[:-1])
      X_with_sens.append(x_with_sens)
      x_control["s1"].append(1 - x1[sensitiveId])
      y.append(int(d['Churn'] == 'Yes')*2 - 1)

  return X_raw, X, X_with_sens, x_control, y, fieldIds

# X_raw, X, X_with_sens, x_control, y, fieldIds = readData("public_data/adult.data")
X_raw, X, X_with_sens, x_control, y, fieldIds = readData("public_data/bank-full.csv")
# X_raw, X, X_with_sens, x_control, y, fieldIds = readData("public_data/default.csv")
# X_raw, X, X_with_sens, x_control, y, fieldIds = readData("public_data/ibm_eval.csv")
# X_raw, X, X_with_sens, x_control, y, fieldIds = readData("public_data/ibm_churn.csv")

frac_to_flip = 0.0

# Apply "synthetic discrimination" using frac_to_flip
for i in range(len(y)):
  if y[i] == 1 and x_control["s1"][i] == 0 and random.random() < frac_to_flip:
    y[i] = -1

X_raw = np.array(X_raw)
X_with_sens = np.array(X_with_sens)
X = np.array(X)
y = np.array(y)
Xt = X_raw.T

# Withhold fraction for testing
testNum = len(X) / 2

x_train = X[:-testNum]
x_test = X[-testNum:]
x_with_sens_train = X_with_sens[:-testNum]
x_with_sens_test = X_with_sens[-testNum:]
y_train = y[:-testNum]
y_test = y[-testNum:]
x_control_train = {}
x_control_test = {}
X_raw_train = X_raw[:-testNum]
X_raw_test = X_raw[-testNum:]
for v in x_control:
  x_control[v] = np.array(x_control[v])
  x_control_train[v] = x_control[v][:-testNum]
  x_control_test[v] = x_control[v][-testNum:]

ut.compute_p_rule(x_control["s1"], y)

apply_fairness_constraints = None
apply_accuracy_constraint = None
sep_constraint = None

loss_function = lf._logistic_loss
sensitive_attrs = ["s1"]
sensitive_attrs_to_cov_thresh = {}
gamma = None


def train_test_classifier(x_train_,x_test_):
  w = ut.train_model(x_train_, y_train, x_control_train, loss_function, apply_fairness_constraints, apply_accuracy_constraint, sep_constraint, sensitive_attrs, sensitive_attrs_to_cov_thresh, gamma)
  print("w = " + str(w))
  train_score, test_score, correct_answers_train, correct_answers_test = ut.check_accuracy(w, x_train_, y_train, x_test_, y_test, None, None)
  distances_boundary_test = (np.dot(x_test_, w)).tolist()
  all_class_labels_assigned_test = np.sign(distances_boundary_test)
  correlation_dict_test = ut.get_correlations(None, None, all_class_labels_assigned_test, x_control_test, sensitive_attrs)
  cov_dict_test = ut.print_covariance_sensitive_attrs(None, x_test_, distances_boundary_test, x_control_test, sensitive_attrs)
  p_rule = ut.print_classifier_fairness_stats([test_score], [correlation_dict_test], [cov_dict_test], sensitive_attrs[0])
  return w, p_rule, test_score

def preds(x, w):
  preds = []
  for x_ in x:
    pred = sigmoid(np.dot(w, x_))
    preds.append(pred)
  return preds

def countAcc(preds, y):
  a = 0
  for pred, y_ in zip(preds, y):
    if (pred > 0.5) == (y_ == 1):
      a += 1
  acc = a * 1.0 / len(y)
  return acc

def countP(x, sens, preds, y, w):
  n_men = sum(sens)
  n_women = len(sens) - n_men
  sigma_men = 0
  sigma_women = 0
  for pred,s in zip(preds, sens):
    if pred > 0.5:
      if s == 1:
        sigma_men += 1
      else:
        sigma_women += 1
  q_men = sigma_men * 1.0 / n_men
  q_women = sigma_women * 1.0 / n_women
  P = 100*q_women/q_men
  acc = countAcc(preds, y)
  return n_men, n_women, sigma_men, sigma_women, q_men, q_women, P, acc

# Naive thresholding scheme (assuming women are the protected class)
def naive(x, sens, y, w, p_to_achieve, acc_to_achieve):
  ps = preds(x, w)
  n_men, n_women, sigma_men, sigma_women, q_men, q_women, P, orig_acc = countP(x, sens, ps, y, w)
  to_flip = []
  for i,r,s,pred in zip(range(len(x)), x, sens, ps):
    if s == 1 and pred > 0.5: # Man
      to_flip.append((p_to_achieve / (100.0 * n_men * (2*pred - 1)), i, -1, pred))
    elif s == 0 and pred < 0.5:
      to_flip.append((1.0 / (n_women * (1 - 2*pred)), i, 1, pred))
  flipped_women = 0
  flipped_men = 0
  to_flip.sort()
  to_flip.reverse()
  y2 = ps[:]
  p_at_same_acc = None
  acc_at_same_p = None
  best_acc_so_far = 0
  for s, i, y_, pred in to_flip:
    y2[i] = y_
    if y_ == 1:
      flipped_women += 1
    else:
      flipped_men += 1
    _, _, _, _, _, _, newP, _ = countP(x, sens, y2, y, w)
    new_acc = countAcc(y2, y)
    if new_acc >= acc_to_achieve:
      p_at_same_acc = (newP, new_acc) # P should be increasing
    if newP >= p_to_achieve and acc_at_same_p == None:
      acc_at_same_p = new_acc
    print s, i, y_, pred, new_acc, newP
    if newP >= 100:
      break
    if new_acc > best_acc_so_far:
      best_acc_so_far = new_acc
  return p_at_same_acc, acc_at_same_p

# scikit logistic regression to predict the sensitive variable
model = LogisticRegression(class_weight='balanced', fit_intercept = False)
model = LogisticRegression()
y_ = [bool(1 - a) for a in x_control_train["s1"]]
y_test_ = [bool(1 - a) for a in x_control_test["s1"]]
model = model.fit(x_train, y_)
w_sensitive = model.coef_[0]

predicted = model.predict(x_test)
probs = model.predict_proba(x_test)

print metrics.accuracy_score(y_test_, predicted)
print metrics.roc_auc_score(y_test_, probs[:, 1])
print metrics.confusion_matrix(y_test_, predicted)
print metrics.classification_report(y_test_, predicted)

random.seed(0)
""" Classify the data while optimizing for accuracy """
print "== Unconstrained (original) classifier =="
# all constraint flags are set to 0 since we want to train an unconstrained (original) classifier
apply_fairness_constraints = 0
apply_accuracy_constraint = 0
sep_constraint = 0
w_uncons, p_uncons, acc_uncons = train_test_classifier(x_train,x_test)

random.seed(0)
print "== Classifier with fairness constraint =="
apply_fairness_constraints = 1 # set this flag to one since we want to optimize accuracy subject to fairness constraints
apply_accuracy_constraint = 0
sep_constraint = 0
sensitive_attrs_to_cov_thresh = {"s1": 0}
w_f_cons, p_f_cons, acc_f_cons  = train_test_classifier(x_train,x_test)

print "Naive accuracy:    ", acc_uncons
print "Naive P:           ", p_uncons
print "Fairness accuracy: ", acc_f_cons
print "Fairness P:        ", p_f_cons

p_at_same_acc, acc_at_same_p = naive(x_test, x_control_test["s1"], y_test, w_uncons, p_f_cons, min(acc_f_cons, acc_uncons))

print "P@constrained accuracy: ", p_at_same_acc
print "accuracy@constrained P: ", acc_at_same_p