# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cvxpy as cp
import seaborn as sns
import emm
import warnings
import sklearn as sk
warnings.filterwarnings('ignore')

# Plot styles
plt.style.use('seaborn-notebook')

import matplotlib.pyplot as plt

SMALL_SIZE = 15
MEDIUM_SIZE = 16
BIGGER_SIZE = 22

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

# +
processed_dir = "../assets/data/processed/pima_diabetes/"
# Get data into dataframe
df = pd.read_csv(processed_dir + "diabetes.csv",index_col=0)

# Replace 0 with nan
nan_cols = ['Glucose', 'BloodPressure','SkinThickness','Insulin','BMI']
df[nan_cols]=df[nan_cols].replace({'0':np.nan,0:np.nan})
# -

fig, ax = plt.subplots(4,2, figsize=(16,16))
sns.distplot(df.Age, bins = 20, ax=ax[0,0]) 
sns.distplot(df.Pregnancies, bins = 20, ax=ax[0,1]) 
sns.distplot(df.Glucose, bins = 20, ax=ax[1,0]) 
sns.distplot(df.BloodPressure, bins = 20, ax=ax[1,1]) 
sns.distplot(df.SkinThickness, bins = 20, ax=ax[2,0])
sns.distplot(df.Insulin, bins = 20, ax=ax[2,1])
sns.distplot(df.DiabetesPedigreeFunction, bins = 20, ax=ax[3,0]) 
sns.distplot(df.BMI, bins = 20, ax=ax[3,1]) 
plt.show()

marginals_mean = df.groupby('Outcome').mean()
m_0 = marginals_mean.loc[0]
m_1 = marginals_mean.loc[1]
display(marginals_mean.T)
print(marginals_mean.T.round(2).to_latex())

df[df['Outcome']==1].shape[0] / df.shape[0]

# +
from emm.losses import LeastSquaresLoss
from emm.losses import CorpusKLLoss
from emm.losses import EqualityLoss
loss_gl0 = CorpusKLLoss(m_0['Glucose'])
loss_gl1 = CorpusKLLoss(m_1['Glucose'])
loss_bmi0 = CorpusKLLoss(m_0['BMI'])
loss_bmi1 =CorpusKLLoss(m_1['BMI'])
loss_age0 = CorpusKLLoss(m_0['Age'])
loss_age1 = CorpusKLLoss(m_1['Age'])


# margs = {0 : [emm.reweighting.marginal('Glucose', loss_gl0.fun, loss_gl0),
#               emm.reweighting.marginal('BMI', loss_gl1.fun, loss_gl1),
#              emm.reweighting.marginal('Age', loss_bmi0.fun, loss_bmi0)],
#          1 : [emm.reweighting.marginal('Glucose', loss_bmi1.fun, loss_bmi1),
#               emm.reweighting.marginal('BMI', loss_age0.fun, loss_age0),
#              emm.reweighting.marginal('Age', loss_age1.fun, loss_age1)]
#         }

            

margs = {0 : [emm.reweighting.marginal('Glucose', 'mean', LeastSquaresLoss(m_0['Glucose']),True),
              emm.reweighting.marginal('BMI', 'mean', LeastSquaresLoss(m_0['BMI']),True),
             emm.reweighting.marginal('Age', 'mean', LeastSquaresLoss(m_0['Age']),True)],
         1 : [emm.reweighting.marginal('Glucose', 'mean', LeastSquaresLoss(m_1['Glucose'], scale=100),True),
              emm.reweighting.marginal('BMI', 'mean', LeastSquaresLoss(m_1['BMI'], scale=100),True),
             emm.reweighting.marginal('Age', 'mean', LeastSquaresLoss(m_1['Age'], scale =100),True)]
        }

# margs = {0 : [emm.reweighting.marginal('Glucose', 'mean', EqualityLoss(m_0['Glucose'])),
#               emm.reweighting.marginal('BMI', 'mean', EqualityLoss(m_0['BMI'])),
#              emm.reweighting.marginal('Age', 'mean', EqualityLoss(m_0['Age']))],
#          1 : [emm.reweighting.marginal('Glucose', 'mean', EqualityLoss(m_1['Glucose'])),
#               emm.reweighting.marginal('BMI', 'mean', EqualityLoss(m_1['BMI'])),
#              emm.reweighting.marginal('Age', 'mean', EqualityLoss(m_1['Age']))]
#         }
            

target = df
corpus = df.drop(columns = 'Outcome')
from sklearn.linear_model import LogisticRegression
param_grid_log =  {"classifier" : [LogisticRegression()],
                   "classifier__C" : np.logspace(-3,3,5)}

from sklearn.tree import DecisionTreeClassifier
param_grid_tree =  {'classifier' : [DecisionTreeClassifier()],
     'classifier__max_depth' : [5,10,20,50]}

from sklearn.ensemble import RandomForestClassifier
param_grid_rf  = {'classifier' : [RandomForestClassifier()],
    'classifier__bootstrap': [True],
    'classifier__max_depth': [5,10,20,50],
    'classifier__max_features' : ['sqrt'],
    'classifier__n_estimators': [200,300,400]
}
param_grid = [param_grid_log, param_grid_tree, param_grid_rf]

from sklearn.impute import SimpleImputer    
pipe = [("imputer", SimpleImputer())]
    
(rw_corpus, js, ml) = emm.metrics.multiple_models(target,corpus,margs, 
                                           param_grid, regularizer = emm.regularizers.EntropyRegularizer(), 
                                           lam=0.01, verbose=True, pipeline_steps = pipe, 
                                           scoring = sk.metrics.roc_auc_score, metrics=[sk.metrics.accuracy_score, sk.metrics.roc_auc_score, sk.metrics.confusion_matrix])

# +
fig, ax = plt.subplots(2,3, figsize=(16,8))


sns.distplot(df[df['Outcome']==0].Glucose, bins = 'auto', ax=ax[0,0]) 
sns.distplot(df[df['Outcome']==0].BMI, bins = 'auto', ax=ax[0,1]) 
sns.distplot(df[df['Outcome']==0].Age, bins = 'auto', ax=ax[0,2])

emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==0].Glucose, rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[0,0])
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==0].BMI, rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[0,1])
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==0].Age, rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[0,2])

sns.distplot(df[df['Outcome']==1].Glucose, bins = 'auto', ax=ax[1,0]) 
sns.distplot(df[df['Outcome']==1].BMI, bins = 'auto', ax=ax[1,1]) 
sns.distplot(df[df['Outcome']==1].Age, bins = 'auto', ax=ax[1,2])

emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==1].Glucose, rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[1,0])
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==1].BMI, rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[1,1])
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==1].Age, rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[1,2])
plt.tight_layout()
plt.show()
# -

ax

# +
fig, ax = plt.subplots( figsize=(10, 12) , nrows=3, ncols=2) 

ax = ax.flatten()



sns.distplot(df[df['Outcome']==0].Glucose, bins = 'auto', ax=ax[0], label='Target', kde_kws={'linewidth':3}) 
sns.distplot(df[df['Outcome']==0].BMI, bins = 'auto', ax=ax[2], label='Target', kde_kws={'linewidth':3}) 
sns.distplot(df[df['Outcome']==0].Age, bins = 'auto', ax=ax[4],label='Target', kde_kws={'linewidth':3})

emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==0].Glucose, rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[0],label='Reweighted corpus', kde_kws={'linewidth':3})
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==0].BMI, rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[2], label='Reweighted corpus', kde_kws={'linewidth':3})
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==0].Age, rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[4],label='Reweighted corpus', kde_kws={'linewidth':3})

ax[4].legend()

sns.distplot(df[df['Outcome']==1].Glucose, bins = 'auto', ax=ax[1], kde_kws={'linewidth':3}) 
sns.distplot(df[df['Outcome']==1].BMI, bins = 'auto', ax=ax[3], kde_kws={'linewidth':3}) 
sns.distplot(df[df['Outcome']==1].Age, bins = 'auto', ax=ax[5], kde_kws={'linewidth':3})

emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==1].Glucose, rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[1], kde_kws={'linewidth':3})
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==1].BMI, rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[3], kde_kws={'linewidth':3})
emm.utils.weighted_hist(rw_corpus[rw_corpus['Outcome']==1].Age, rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[5], kde_kws={'linewidth':3})    

fig.set_facecolor('w')
#plt.suptitle('Marginal feature distribution for target and reweighted corpus')
ax[0].set_title('Label 0', fontsize=24)
ax[1].set_title('Label 1', fontsize=24)
plt.tight_layout()
plt.savefig('PIMA_dist_reweighted.jpg')
plt.show()
# -

ml

js

print(pd.DataFrame(js).to_latex())

m = len(corpus)
base_test = pd.concat([corpus.copy(), corpus.copy()], ignore_index=True)
base_test['weights'] = np.ones(2*m) /(2*m)
base_test['Outcome'] = np.concatenate([np.zeros(m), np.ones(m)],axis=0)
emm.compute_js_distance(target,base_test)

# +
fig, ax = plt.subplots(figsize=(10, 12) , nrows=3, ncols=2, sharey=True) 
ax = ax.flatten()
    

sns.ecdfplot(df[df['Outcome']==0].Glucose, ax=ax[0], label='Target') 
sns.ecdfplot(df[df['Outcome']==0].BMI, ax=ax[2], label='Target') 
sns.ecdfplot(df[df['Outcome']==0].Age, ax=ax[4], label='Target')

sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==0].Glucose, weights=rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[0], label='Reweighted corpus')
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==0].BMI, weights=rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[2], label='Reweighted corpus')
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==0].Age, weights=rw_corpus[rw_corpus['Outcome']==0].weights, ax=ax[4], label='Reweighted corpus')

sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==0].Glucose, ax=ax[0], label='Unweighted corpus')
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==0].BMI, ax=ax[2], label='Unweighted corpus')
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==0].Age, ax=ax[4], label='Unweighted corpus')

ax[0].legend()

sns.ecdfplot(df[df['Outcome']==1].Glucose, ax=ax[1]) 
sns.ecdfplot(df[df['Outcome']==1].BMI, ax=ax[3]) 
sns.ecdfplot(df[df['Outcome']==1].Age, ax=ax[5])

sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==1].Glucose, weights=rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[1])
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==1].BMI, weights=rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[3])
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==1].Age, weights=rw_corpus[rw_corpus['Outcome']==1].weights, ax=ax[5])    

sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==1].Glucose, ax=ax[1])
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==1].BMI, ax=ax[3])
sns.ecdfplot(x=rw_corpus[rw_corpus['Outcome']==1].Age, ax=ax[5])  

fig.set_facecolor('w')
#plt.suptitle('Conditional feature CDF for target, reweighted corpus, unweighted corpus')

plt.tight_layout()
plt.savefig('PIMA_cdf_reweighted.jpg')
plt.show()
# -

print(emm.utils.weighted_mean(rw_corpus.drop(columns=['Outcome', 'weights']), rw_corpus['weights']))

target['Outcome'].value_counts()

corpus.mean()

base_test = pd.concat([corpus.copy(), corpus.copy()], ignore_index=True)
base_test['weights'] = np.ones(2*m) /(2*m)
base_test['Outcome'] = np.concatenate([np.zeros(m), np.ones(m)],axis=0)
emm.compute_js_distance(target,base_test)

marginals_mean



StarcoderdataJupyterScriptsDedupFiltered