####################
## load libraries ##
####################
import numpy as np
import pandas as pd
np.random.seed(123456789)
from sklearn.tree import DecisionTreeRegressor
from sklearn.tree import DecisionTreeClassifier

############################
## Toy Dataset Simulation ##
############################
def simulate_df(n=100, seed=123456, binary_flag=False):
    np.random.seed(seed)
    
    ## specify dataframe
    df = pd.DataFrame()

    ## 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])
    
    theta_0 = 5.5
    theta_1 = 1.28
    theta_2 = 0.42
    theta_3 = 2.32
    theta_4 = -3.15
    theta_5 = 3.12
    theta_6 = -4.29
    theta_7 = -1.23
    theta_8 = -10.18
    theta_9 = 2.21
    theta_10 = 10.3
    
    if(binary_flag):
        Z = theta_0 + (theta_1*df['L1']) + (theta_2*df['L2']) + (theta_3*df['L3']) + (theta_4*df['L4']) + (theta_5*df['L5']) + (theta_6*df['L6']) + (theta_7*df['L2']*df['L4']) + (theta_8*df['L3']*df['L6']) + (theta_9*df['L5']*df['L5']) + (theta_10*np.sin(df['L5']))
        p = 1 / (1 + np.exp(-Z))
        df['Y'] = np.random.binomial(1, p)
        df.loc[df['Y']==0, 'Y'] = -1
    else:
        df['Y'] = theta_0 + (theta_1*df['L1']) + (theta_2*df['L2']) + (theta_3*df['L3']) + (theta_4*df['L4']) + (theta_5*df['L5']) + (theta_6*df['L6']) + (theta_7*df['L2']*df['L4']) + (theta_8*df['L3']*df['L6']) + (theta_9*df['L5']*df['L5']) + (theta_10*np.sin(df['L5'])) + np.random.normal(0, 0.1, df.shape[0])

    return(df)

###############################################
## Gradient Boosting with Continuous Outcome ##
###############################################

## simulate the dataset with continuous outcome Y
df = simulate_df(n=1000, seed=123456)
df['Y_original'] = df['Y']

## fit Gradient Boosting Model with depth-3 Regression Trees as Weak Learners
## continue fitting model until loss function < 1
alpha = 100000
current_loss = sum((df['Y'])**2)
while(current_loss > 1):
    model = DecisionTreeRegressor(random_state=0, max_depth=3)
    model.fit(df[['L1', 'L2', 'L3','L4', 'L5', 'L6']], df['Y'])
    df['Y_hat'] = model.predict(df[['L1', 'L2', 'L3','L4', 'L5', 'L6']])
    df['Y_hat_squared'] = df['Y_hat']**2
    df['Y_hat_scaled'] = np.sqrt(df['Y_hat_squared'] / df['Y_hat_squared'].sum()) * np.sign(df['Y_hat'])
    loss_not_lowered_flag = True
    while(loss_not_lowered_flag):
        new_loss = sum((df['Y'] - (alpha*df['Y_hat_scaled']))**2)
        if(new_loss < current_loss):
            loss_not_lowered_flag = False
            current_loss = new_loss
            print('Current Loss: ' + str(current_loss))
            df['Y'] = df['Y'] - (alpha*df['Y_hat_scaled'])
        else:
            alpha = 0.99*alpha
    del model
print('model converged')

######################################################
## AdaBoost with Binary Categorical Outcome {-1, 1} ##
######################################################

## simulate the dataset with continuous outcome Y
df = simulate_df(n=1000, seed=123456, binary_flag=True)
df = df.rename(columns={'Y':'Y_original'})
df['w'] = df.shape[0]*[1/df.shape[0]]
df['Y'] = 0

## fit AdaBoost Model with depth-2 Decision Trees as Weak Learners
## continue fitting model until all n-1000 observations predicted correctly
count = True
while(count):
    model = DecisionTreeClassifier(random_state=0, max_depth=2)
    model.fit(df[['L1', 'L2', 'L3','L4', 'L5', 'L6']], df['Y_original'], sample_weight=df['w'])
    df['Y_hat'] = model.predict(df[['L1', 'L2', 'L3','L4', 'L5', 'L6']])
    df.loc[df['Y_hat']==df['Y_original'], 'w'] = 0
    epsilon = sum(df['w'])
    alpha = 0.5*(np.log((1-epsilon)/epsilon))
    
    df['Y'] = df['Y'] + (alpha*df['Y_hat'])
    
    current_loss = sum(np.exp(-df['Y_original']*df['Y']))
    psi = np.exp(-df['Y_original']*df['Y'])
    df['w'] = psi / current_loss
    
    print('Current Loss: ' + str(current_loss))
    
    df['Y_final'] = 1
    df.loc[df['Y']<0, 'Y_final'] = -1
    
    if(df.loc[df['Y_original']!=df['Y_final'], :].shape[0] == 0):
        print('results converged, all datapoints correctly classified')
        break

pd.crosstab(df['Y_original'], df['Y_final'])