####################
## load libraries ##
####################
import numpy as np
import pandas as pd
np.random.seed(12345)
import statsmodels.api as sm
import statsmodels.formula.api as smf

##########################################################################
## simulate dataset and analysis (with and witout covariate adjustment) ##
##########################################################################
def simulate_df(n=200, outcome_type='normal', A_on_Y=2, data_only=False, seed=123456):
    np.random.seed(seed)
    ## specify dataframe
    df = pd.DataFrame()
    df['A'] = np.random.choice(int(n/2)*[0] + int(n/2)*[1], size=n, replace=False)

    ## specify variables L1 through L6
    L1_split = 0.52
    L2_split = 0.23
    L3_split = 0.38
    df['L1'] = np.random.choice([0, 1], size=n, replace=True, p=[L1_split, (1-L1_split)])
    df['L2'] = np.random.choice([0, 1], size=n, replace=True, p=[L2_split, (1-L2_split)])
    df['L3'] = np.random.choice([0, 1], size=n, replace=True, p=[L3_split, (1-L3_split)])
    df['L4'] = np.random.normal(0, 1, df.shape[0])
    df['L5'] = np.random.normal(0, 0.75, df.shape[0])
    df['L6'] = np.random.normal(0, 2, df.shape[0])
    
    df_A0 = df.copy()
    df_A0['A'] = 0
    df_A1 = df.copy()
    df_A1['A'] = 1
    
    theta_0 = -0.5
    theta_1 = A_on_Y
    theta_2 = 0.28
    theta_3 = 0.42
    theta_4 = 0.32
    theta_5 = -0.15
    theta_6 = 0.12
    theta_7 = -0.29
    
    if(outcome_type=='normal'):
        df['Y'] = theta_0 + (theta_1*df['A']) + (theta_2*df['L1']) + (theta_3*df['L2']) + (theta_4*df['L3']) + (theta_5*df['L4']) + (theta_6*df['L5']) + (theta_7*df['L6']) + np.random.normal(0, 1, df.shape[0])
        if(data_only):
            return(df)
        model_unadjusted = smf.ols('Y ~ A', data=df).fit()
        print(model_unadjusted.summary()) 
        model_adjusted = smf.ols('Y ~ A + L1 + L2 + L3 + L4 + L5 + L6', data=df).fit()
        print(model_adjusted.summary())
        print('The sample variance is: ' + str(df['Y'].var()))
        
    if(outcome_type=='binomial'):
        df['Z'] = theta_0 + (theta_1*df['A']) + (theta_2*df['L1']) + (theta_3*df['L2']) + (theta_4*df['L3']) + (theta_5*df['L4']) + (theta_6*df['L5']) + (theta_7*df['L6'])
        df['p'] = 1 / (1 + np.exp(-df['Z']))
        df['Y'] = np.random.binomial(1, df['p'])
        if(data_only):
            return(df)
        model_unadjusted = smf.glm('Y ~ A', data=df, family=sm.families.Binomial()).fit()
        print(model_unadjusted.summary()) 
        model_adjusted = smf.glm('Y ~ A + L1 + L2 + L3 + L4 + L5 + L6', data=df, family=sm.families.Binomial()).fit()
        print(model_adjusted.summary())
    
    
    df_A0['Y_pred_unadjusted'] =  model_unadjusted.predict(df_A0)
    df_A1['Y_pred_unadjusted'] =  model_unadjusted.predict(df_A1)
    df_A0['Y_pred_adjusted'] =  model_adjusted.predict(df_A0)
    df_A1['Y_pred_adjusted'] =  model_adjusted.predict(df_A1)
       
    print('The unadjusted Causal Effect Difference is: ' + str(df_A1['Y_pred_unadjusted'].mean() - df_A0['Y_pred_unadjusted'].mean()))
    print('The adjusted Causal Effect Difference is: ' + str(df_A1['Y_pred_adjusted'].mean() - df_A0['Y_pred_adjusted'].mean()))   
    
    return(df)

######################
## Power Simulation ##
######################
def power_analysis(df_original, outcome_type='normal', m=10, alpha=0.05, A_on_Y=0.5, seed=123456):
    np.random.seed(seed)
    
    df_p_values = pd.DataFrame()
    df_p_values['m'] = list(range(0,m))
    df_p_values['p_unadjusted'] = None
    df_p_values['p_adjusted'] = None
    
    for i in range(0, m):
        #df = df_original.sample(n=df_original.shape[0], replace=True, axis=0).reset_index(drop=True)

        if(outcome_type=='normal'):
            df = simulate_df(n=df_original.shape[0], outcome_type='normal', A_on_Y=A_on_Y, data_only=True, seed=i)            
            model_unadjusted = smf.ols('Y ~ A', data=df).fit()
            df_p_values.loc[i, 'p_unadjusted'] = model_unadjusted.pvalues['A']
            model_adjusted = smf.ols('Y ~ A + L1 + L2 + L3 + L4 + L5 + L6', data=df).fit()
            df_p_values.loc[i, 'p_adjusted'] = model_adjusted.pvalues['A']
            del df, model_unadjusted, model_adjusted
            
        if(outcome_type=='binomial'):
            df = simulate_df(n=df_original.shape[0], outcome_type='binomial', A_on_Y=A_on_Y, data_only=True, seed=i) 
            model_unadjusted = smf.glm('Y ~ A', data=df, family=sm.families.Binomial()).fit()
            df_p_values.loc[i, 'p_unadjusted'] = model_unadjusted.pvalues['A']
            model_adjusted = smf.glm('Y ~ A + L1 + L2 + L3 + L4 + L5 + L6', data=df, family=sm.families.Binomial()).fit()
            df_p_values.loc[i, 'p_adjusted'] = model_adjusted.pvalues['A']
            del df, model_unadjusted, model_adjusted
    
    unadjusted_power = 1 - (df_p_values.loc[df_p_values['p_unadjusted']>alpha, :].shape[0] / m)
    adjusted_power = 1 - (df_p_values.loc[df_p_values['p_adjusted']>alpha, :].shape[0] / m)
    
    print('The unadjusted power is: ' + str(unadjusted_power))
    print('The adjusted power is: ' + str(adjusted_power))

#############################
## Continuous Outcome Case ##
#############################
## Recover unbiased estimate of Mean Causal Effect Difference, with and without covariate adjustment:
df_normal = simulate_df(n=134, outcome_type='normal', A_on_Y=0.5, data_only=False, seed=1081565)
## conduct power analysis, both with and without covariate adjustment:
power_analysis(df_normal, outcome_type='normal', m=10000, alpha=0.05, A_on_Y=0.5, seed=123)

#########################
## Binary Outcome Case ##
#########################
## Recover unbiased estimate of Mean Causal Effect Difference, with and without covariate adjustment:
df_binomial = simulate_df(n=652, outcome_type='binomial', A_on_Y=0.5, data_only=False, seed=12345678)
## conduct power analysis, both with and without covariate adjustment:
power_analysis(df_binomial, outcome_type='binomial', m=10000, alpha=0.05, A_on_Y=0.5, seed=123456)