import datasethelper,lewishelper
import copy
import pandas as pd
import sys

model=sys.argv[1]
col='score_text'
#col='is_recid'
(X_train,y_train,df,pred,clf)=datasethelper.process_compas(col,model)
#df=datasethelper.process_compas('is_recid')
#df[col]=pred

featlst=list(df.columns)

featlst.remove(col)


if col=='score_text' and 'is_recid' in featlst:
    featlst.remove('is_recid')
def get_majority(dic1,s1,s2,feat):
    score=[0,0,0]
    if feat in dic1.keys():
        score=dic1[feat]
        
    i=0
    while i<3:
        score[i]=max(score[i],s1[i])
        score[i]=max(score[i],s2[i])
        i+=1
    dic1[feat]=score
    return dic1

def get_majority(dic1,s1,s2,feat):
    score=[0,0,0]
    if feat in dic1.keys():
        score=dic1[feat]

    i=0
    while i<3:
        score[i]=max(score[i],s1[i])
        score[i]=max(score[i],s2[i])
        i+=1
    dic1[feat]=score
    return dic1


score_dic={}
for feat in featlst:
    #print (feat)
    uniqval=list(set(list(df[feat])))

    if feat=='age_cat' or feat=='sex' or feat=='race':
        i=0
        while i<len(uniqval):
            j=i+1
            while j<len(uniqval):
                score1=lewishelper.get_scores_regression(df,[feat],[uniqval[i]],[uniqval[j]],[],[],[],col)
                score2=lewishelper.get_scores_regression(df,[feat],[uniqval[j]],[uniqval[i]],[],[],[],col)
                score_dic=get_majority(score_dic,score1,score2,feat)
                j+=1
            i+=1
    else:
        i=0
        while i<len(uniqval):
            j=i+1
            while j<len(uniqval):
                score1=lewishelper.get_scores_regression(df,[feat],[uniqval[i]],[uniqval[j]],[],[],['age_cat','sex','race'],col)
                score2=lewishelper.get_scores_regression(df,[feat],[uniqval[j]],[uniqval[i]],[],[],['age_cat','sex','race'],col)
                score_dic=get_majority(score_dic,score1,score2,feat)
                j+=1
            i+=1

#print (score_dic)



n_score={}
s_score={}
sn_score={}
for feat in score_dic.keys():
    n_score[feat]=score_dic[feat][0]
    s_score[feat]=score_dic[feat][1]
    sn_score[feat]=score_dic[feat][2]
    #print (feat,n_score[feat],s_score[feat],sn_score[feat])


sorted_s = sorted(s_score.items(), key=lambda kv: kv[1],reverse=True)
sorted_n = sorted(n_score.items(), key=lambda kv: kv[1],reverse=True)
sorted_sn = sorted(sn_score.items(), key=lambda kv: kv[1],reverse=True)

slst=[]
nlst=[]
snlst=[]
svallst=[]
nvallst=[]
snvallst=[]

i=0
for (x,y) in sorted_s:
    slst.append(x)
    svallst.append(y)
    nlst.append(sorted_n[i][0])
    nvallst.append(sorted_n[i][1])
    snlst.append(sorted_sn[i][0])
    snvallst.append(sorted_sn[i][1])
    i+=1

#print ("sn")
#print (snvallst)
#print (snlst)
# In[ ]:

#print("n")
#print (nlst)
#print (nvallst)

#print("S")
#print (slst)
#print(svallst)





sn=snvallst
n=nvallst
s=svallst
snname=snlst
nname=nlst
sname=slst


import pandas as pd
from functools import reduce

sn = pd.DataFrame(list(zip(sn, snname)), columns=['sn', 'name'])
n = pd.DataFrame(list(zip(n, nname)), columns=['n', 'name'])
s = pd.DataFrame(list(zip(s, sname)), columns=['s', 'name'])
dfs = [sn, n, s]
allScores = reduce(lambda left,right: pd.merge(left,right,on='name'), dfs) #join

allScores = allScores.sort_values(by='sn', ascending=False)
import scipy.stats as ss
nrank=ss.rankdata(allScores['n'])
srank=ss.rankdata(allScores['s'])
snrank=ss.rankdata(allScores['sn'])

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pylab as plot

fsize=20


params = {'legend.fontsize': fsize/1.5,
          'legend.handlelength': 2}
plot.rcParams.update(params)

font = {'family' : "sans serif", 'size'   : fsize}
matplotlib.rcParams['font.family'] = "sans-serif"
matplotlib.rc('font', **font)

labels = allScores['name']
x = np.arange(len(labels))  # the label locations
width = 0.25  # the width of the bars

# plt.figure(figsize=(6, 10)) # in inches!

fig, ax = plt.subplots()
rects1 = ax.barh(x - width, allScores['n'], width, label='Nec', color='cornflowerblue', edgecolor='black', hatch="//")
rects2 = ax.barh(x, allScores['s'], width, label='Suf', color='gold', edgecolor='black', hatch="\\\\")
rects3 = ax.barh(x + width, allScores['sn'], width, label='NeSuf', color='forestgreen', hatch="|")

plt.gca().invert_yaxis()
# Add some text for labels, title and custom x-axis tick labels, etc.
ax.set_xlabel('Scores',labelpad=fsize/2)
ax.set_yticks(x)
ax.set_yticklabels(labels)
ax.legend()
# plt.legend(bbox_to_anchor=(.2, 1.2, 1., .102), loc=3,
#        ncol=3, borderaxespad=0., fontsize=fsize)
ax.xaxis.set_ticks_position('top')
ax.xaxis.set_label_position('top')
plt.xticks(np.arange(0, 1.5, .5))
figure = plt.gcf() # get current figure
figure.set_size_inches(5,7.5)

def autolabel(rects,rank):
    """Attach a text label above each bar in *rects*, displaying its height."""
    i=0
    for rect in rects:
        height = rect.get_width()
        val=len(labels)-int(rank[i])+1
        ax.annotate('{}'.format(val),
                    xy=(height - 0.04, rect.get_y() + 0.25),
                    xytext=(15, 0),  # 3 points vertical offset
                    textcoords="offset points",
                    ha='center', va='bottom', fontsize=fsize/1.5)
        i+=1


autolabel(rects1,nrank)
autolabel(rects2,srank)
autolabel(rects3,snrank)

matplotlib.rcParams['hatch.linewidth'] = 0.3

fig.tight_layout()

plt.savefig('newplots/compas_globalExplanations_score_text.pdf')
