import os
print (os.environ['CONDA_DEFAULT_ENV'])

from platform import python_version

import multiprocessing
!python --version
import platform
print(platform.platform())
print("cpu cores: {0}".format(multiprocessing.cpu_count()))

# Import dependencies
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import math
from sklearn.linear_model import LogisticRegression

sns.set_style('darkgrid')

df = pd.read_csv('AncillaryScoring_insurance.csv',encoding='latin-1')
df = df.sample(n=8000).reset_index()
df['index'] = df.index

df.head(3)

y = df[['INS_FLAG']]
df = df.drop(columns = ['INS_FLAG'])

try:
    df['treatment'] = df.SEAT_CATEGORY
    df = df.drop(columns = ['SEAT_CATEGORY'], axis=1)
except:
    print()

df.info()

var = 'BAGGAGE_CATEGORY'
pd.pivot_table(df[['treatment',var,'Id']], \
               values = 'Id', index = 'treatment', columns = var, aggfunc= np.count_nonzero)

df.corr()

try:
    df = df.drop(columns = ['PAXCOUNT','PURCHASELEAD','LENGTHOFSTAY','flight_hour'], axis=1)
    df = df.drop(columns = ['ROUTE'], axis=1) # reduce compute resources
except:
    print()

df.info()
df.sample()

T = df.treatment
X = df.loc[:,df.columns !='treatment']
X.shape
df.shape
T.shape

X_encoded = pd.get_dummies(X, columns = ['SALESCHANNEL', 'TRIPTYPEDESC','flight_day','geoNetwork_country'], \
                           prefix = {'SALESCHANNEL':'sc', 'TRIPTYPEDESC':'tt','flight_day':'fd','geoNetwork_country':'geo'}, drop_first=False)

from sklearn.linear_model import LogisticRegression as lr

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import NearestNeighbors
from sklearn import metrics

# Design pipeline to build the treatment estimator
pipe = Pipeline([
    ('scaler', StandardScaler()),
    ('logistic_classifier', lr())
])

pipe.fit(X_encoded, T)

predictions = pipe.predict_proba(X_encoded)
predictions_binary = pipe.predict(X_encoded)

print('Accuracy: {:.4f}\n'.format(metrics.accuracy_score(T, predictions_binary)))
print('Confusion matrix:\n{}\n'.format(metrics.confusion_matrix(T, predictions_binary)))
print('F1 score is: {:.4f}'.format(metrics.f1_score(T, predictions_binary)))

def logit(p):
    logit_value = math.log(p / (1-p))
    return logit_value

predictions_logit = np.array([logit(xi) for xi in predictions[:,1]])

# Density distribution of propensity score (logic) broken down by treatment status

sns.set(rc={"figure.figsize":(14, 6)}) #width=8, height=4
fig, ax = plt.subplots(1,2)
fig.suptitle('Density distribution plots for propensity score and logit(propensity score).')
sns.kdeplot(x = predictions[:,1], hue = T , ax = ax[0])
ax[0].set_title('Propensity Score')
sns.kdeplot(x = predictions_logit, hue = T , ax = ax[1])
ax[1].axvline(-0, ls='--')
ax[1].set_title('Logit of Propensity Score')
#ax[1].set_xlim(-5,5)
plt.show()

print(max(predictions_logit),min(predictions_logit))

print(max(predictions_logit),min(predictions_logit))
common_support = (predictions_logit > -15) & (predictions_logit < 10)

df.loc[:,'propensity_score'] = predictions[:,1]
df.loc[:,'propensity_score_logit'] = predictions_logit
df.loc[:,'outcome'] = y.INS_FLAG

X_encoded.loc[:,'propensity_score'] = predictions[:,1]
X_encoded.loc[:,'propensity_score_logit'] = predictions_logit
X_encoded.loc[:,'outcome'] = y.INS_FLAG
X_encoded.loc[:,'treatment'] = df.treatment

caliper = np.std(df.propensity_score) * 0.25

print('\nCaliper (radius) is: {:.4f}\n'.format(caliper))

df_data = X_encoded

knn = NearestNeighbors(n_neighbors=10 , p = 2, radius=caliper)
knn.fit(df_data[['propensity_score_logit']].to_numpy())

# Common support distances and indexes
distances , indexes = knn.kneighbors(
    df_data[['propensity_score_logit']].to_numpy(), \
    n_neighbors=10)

print('For item 0, the 4 closest neighbors are (ignore first):')
for ds, idx  in zip(distances[0,0:4], indexes[0,0:4]):
    print('item index: {}'.format(idx))
    print('item distance: {:4f}'.format(ds))
print('...')

def propensity_score_matching(row, indexes, df_data):
    current_index = int(row['index']) # index column
    prop_score_logit = row['propensity_score_logit']
    for idx in indexes[current_index,:]:
        if (current_index != idx) and (row.treatment == 1) and (df_data.loc[idx].treatment == 0):
            return int(idx)
         
df_data['matched_element'] = df_data.reset_index().apply(propensity_score_matching, axis = 1, args = (indexes, df_data))

treated_with_match = ~df_data.matched_element.isna()

treated_matched_data = df_data[treated_with_match][df_data.columns]
treated_matched_data.head(3)

cols = treated_matched_data.columns[0:-1]

def prop_matches(row, data, cols):
    return data.loc[row.matched_element][cols]

%%time

untreated_matched_data = pd.DataFrame(data = treated_matched_data.matched_element)

"""attributes = ['Age', 'SibSp', 'Parch', 'Fare', 'sex_female', 'sex_male', 'embarked_C',
       'embarked_Q', 'embarked_S', 'class_1', 'class_2', 'class_3',
       'propensity_score', 'propensity_score_logit', 'outcome', 'treatment']
"""
for attr in cols:
    untreated_matched_data[attr] = untreated_matched_data.apply(prop_matches, axis = 1, data = df_data, cols = attr)
    
untreated_matched_data = untreated_matched_data.set_index('matched_element')
untreated_matched_data.head(3)

untreated_matched_data.shape

treated_matched_data.shape

all_mached_data = pd.concat([treated_matched_data, untreated_matched_data]).reset_index()

all_mached_data.index.is_unique

all_mached_data.treatment.value_counts()

sns.set(rc={"figure.figsize":(14, 8)})
fig, ax = plt.subplots(2,1)
fig.suptitle('Comparison of {} split by outcome and treatment status'.format('propensity_score_logit'))
sns.stripplot(data = df_data.sample(n=1000), y = 'outcome', x = 'propensity_score_logit', 
              #alpha=0.7,
              hue = 'treatment', orient = 'h', ax = ax[0]).set(title = 'Before matching', xlim=(-4, 2))
sns.stripplot(data = all_mached_data.sample(n=1000), y = 'outcome', x = 'propensity_score_logit', 
              #alpha=0.7,
              hue = 'treatment', ax = ax[1] , orient = 'h').set(title = 'After matching', xlim=(-4, 2))
plt.subplots_adjust(hspace = 0.6)
plt.show()

args = ['FNB_CATEGORY','BAGGAGE_CATEGORY','propensity_score_logit']

def plot(arg):
    fig, ax = plt.subplots(1,2)
    fig.suptitle('Comparison of {} split by treatment status.'.format(arg))
    sns.kdeplot(data = df_data, x = arg, hue = 'treatment', ax = ax[0]).set(title='Density before matching')
    sns.kdeplot(data = all_mached_data, x = arg, hue = 'treatment', ax = ax[1]).set(title='Density after matching')
    plt.show()

for arg in args:
    plot(arg)

def cohenD (tmp, metricName):
    treated_metric = tmp[tmp.treatment == 1][metricName]
    untreated_metric = tmp[tmp.treatment == 0][metricName]
    
    d = ( treated_metric.mean() - untreated_metric.mean() ) / math.sqrt(((treated_metric.count()-1)*treated_metric.std()**2 + (untreated_metric.count()-1)*untreated_metric.std()**2) / (treated_metric.count() + untreated_metric.count()-2))
    return d

%%time

data = []
cols = ['BAGGAGE_CATEGORY', 'FNB_CATEGORY',
       'flightDuration_hour', 'sc_Internet', 'sc_Mobile', 'tt_CircleTrip',
       'tt_OneWay', 'tt_RoundTrip', 'fd_Fri', 'fd_Mon', 'fd_Sat', 'fd_Sun',
       'fd_Thu', 'fd_Tue', 'fd_Wed', 'geo_Vietnam']
for cl in cols:
    data.append([cl,'before', cohenD(df_data,cl)])
    data.append([cl,'after', cohenD(all_mached_data,cl)])

res = pd.DataFrame(data, columns=['variable','matching','effect_size'])

sn_plot = sns.barplot(data = res, y = 'variable', x = 'effect_size', hue = 'matching', orient='h')
sn_plot.set(title='Standardised Mean differences accross covariates before and after matching')

cols.append('treatment')

overview = all_mached_data[['outcome','treatment']].groupby(by = ['treatment']).aggregate([np.mean, np.var, np.std, 'count'])
print(overview)

treated_outcome = overview['outcome']['mean'][1]
treated_counterfactual_outcome = overview['outcome']['mean'][0]

att = treated_outcome - treated_counterfactual_outcome
print('The Average Treatment Effect (ATT): {:.4f}'.format(att))

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

from causalinference import CausalModel

# y observed outcome, 
# T treatment indicator, 
# X covariate matrix (not one hot encoded),
causal = CausalModel(y.values, T, X)
print(causal.summary_stats)