#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
FUNCTION: At the top, include a short description of the function of the script.
		  That is, in one sentence answer the question, 'What does this script do?'.
		  Also provide quick note on any bugs or issues as well as last time that tests were conducted.

TESTING:
    Number conducted = 0
    Number passed = 0
    Pass percentageon last 10 = 0

LAST UPDATED: DD/MM/YYYY

OWNER: Your Name
'''

# Packages and modules
# System
import os
import sys

# Data Manipulation
import pandas as pd
import numpy as np

# Visualizations
import seaborn as sns
import matplotlib.pyplot as plt

# Analytics
import researchpy as rp
import statsmodels.api as sm
import statsmodels.formula.api as smf
import scipy.stats as stats
from scipy.optimize import curve_fit
from sklearn.ensemble import RandomForestRegressor as RFR
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import AdaBoostRegressor

import random
import warnings
warnings.filterwarnings("ignore")


# Script preference
pd.set_option("display.max_columns", None)

#%% Read in and prep data
data = pd.read_csv('../Data/all_indifference_props.csv')

# Transform to long df
df_long = pd.melt(data, id_vars=['Group', 'Commodity', 'Gain_Loss','Amount', 'Delay_Prob', 'IV'])
df_long = df_long.rename(columns={'variable':'pid', 'value':'indiff_point'})
df_long.to_csv('../Data/data_long.csv')
df_long

#%% Mixed Effects Model Linear Regression
# Take a look at the data to see if anything looks funky
rp.codebook(df_long)

#%% Get model
model = smf.mixedlm(formula="indiff_point ~ C(Commodity)+ C(Gain_Loss) + C(Amount) + C(Delay_Prob) + IV",
                    data=df_long, groups=df_long["Group"], missing='drop', re_formula='~IV').fit()
model.summary()

#%% One hot encode categorical variables for predictions
encode_df = pd.DataFrame(df_long['indiff_point'])
encode_df.columns = ['indiff_point']
encode_df = pd.concat([encode_df, pd.get_dummies(df_long['Group'], prefix='Group')], axis=1)
encode_df = pd.concat([encode_df, pd.get_dummies(df_long['Commodity'], prefix='Commodity')], axis=1)
encode_df = pd.concat([encode_df, pd.get_dummies(df_long['Gain_Loss'], prefix='Gain_Loss')], axis=1)
encode_df = pd.concat([encode_df, pd.get_dummies(df_long['Delay_Prob'], prefix='Delay_Prob')], axis=1)
encode_df['Amount'] = df_long['Amount']
encode_df['IV'] = df_long['IV']
encode_df['pid'] = df_long['pid']
encode_df = encode_df.dropna()
encode_df = encode_df.reset_index(drop=True)

encode_df = encode_df.drop(['Group_Control', 'Gain_Loss_Loss', 'Delay_Prob_Probability'], axis=1)

encode_df.head(20)

#%% Functions
def plot_pred(predictors, observed_data, title, save_name=None, rand_state=5687463,
              model='RF'):
    # Fit the model
    if model=='RF':
        regr = RFR(max_depth=5, random_state=rand_state)
    elif model=='LR':
        regr = LinearRegression()
    elif model=='GBR':
        regr = GradientBoostingRegressor(random_state=rand_state)
    elif model=='ABR':
        regr = AdaBoostRegressor(random_state=rand_state)

    regr.fit(predictors, observed_data)
    r2 = regr.score(predictors, observed_data)

    # Plot it
    fig, ax = plt.subplots(figsize=(10, 10))
    sns.regplot(x=observed_data, y=regr.predict(predictors), color='k', )
    plt.title(title, fontsize=36, pad=16)
    ax.text(0.7, 0.08, r"$r^2=$"+f'{round(r2, 2)}', fontsize=24)
    plt.xticks(fontsize=20)
    plt.xlim(0, 1)
    plt.yticks(fontsize=20)
    plt.ylim(0, 1)
    plt.ylabel("Predicted Indifference Proportion", fontsize=30, labelpad=20)
    plt.xlabel("Observed Indifference Proportion", fontsize=30, labelpad=20)
    # plt.savefig(f'../Figures/{save_name}.png', bbox_inches="tight")
    plt.show()

    return r2

#%% Lists to store vals for later
fit = []
model = []
vac = []

for model in ['LR', 'RF', 'GBR', 'ABR']:
    # Predictor cols
    preds = ['Group_Cocaine', 'Commodity_Cocaine', 'Commodity_Health', 'Commodity_Money',
             'Gain_Loss_Gains', 'Delay_Prob_Delay', 'Amount', 'IV']
    
    for i in range(0, 1000):
        rand_state = int(round(random.random()*10000, 0))
    
        # Fit to all data
        fit.append('all')
        vac.append(plot_pred(predictors=encode_df[preds].astype(float),
                  observed_data=encode_df[['indiff_point']].astype(float).to_numpy(),
                  title='RF: All Participants', save_name='reg_plot_all_participants',
                   rand_state=rand_state, model=model))
        
        # Cocaine group only
        fit.append('cocaine')
        temp_df = encode_df[encode_df['Group_Cocaine']==1]
        vac.append(plot_pred(predictors=temp_df[preds].astype(float),
                  observed_data=temp_df[['indiff_point']].astype(float).to_numpy(),
                  title='RF: All Participants', save_name='reg_plot_cocaine',
                   rand_state=rand_state, model=model))
        
        # Control group only
        fit.append('control')
        temp_df = encode_df[encode_df['Group_Cocaine']==0]
        vac.append(plot_pred(predictors=temp_df[preds].astype(float),
                  observed_data=temp_df[['indiff_point']].astype(float).to_numpy(),
                  title='RF: All Participants', save_name='reg_plot_control',
                   rand_state=rand_state, model=model))
        
        # Fit to individual participant data
        for p in encode_df['pid'].unique():
            fit.append('ind')
            temp_df = encode_df[encode_df['pid']==p].dropna()
            temp_df = temp_df.reset_index(drop=True)
            val = plot_pred(predictors=temp_df[['Gain_Loss_Gains', 'Delay_Prob_Delay', 'Amount', 'IV']].astype(float),
                            observed_data=temp_df[['indiff_point']].astype(float).to_numpy(),
                            title=f'RF: Participant {p}', save_name=f'reg_plot_random_forest_{p}',
                            rand_state=rand_state, model=model)
            vac.append(val)

        print(f'Using {model}, {(i+1)} iterations complete ({round(100*(i/1000), 2)}%)')


# Make a df
fit_df = pd.DataFrame({'fit':fit, 'vac':vac})
fit_df.to_csv('./Data/boostrapped_ml_regressors.csv')

#%% plot it
fig, ax = plt.subplots(figsize=(12, 8))
sns.histplot(data=fit_df, x='vac', hue='fit', bins=22)
plt.xticks(fontsize=20)
plt.xlim(0, 1)
plt.yticks(fontsize=20)
plt.ylim(0, )
plt.ylabel("Number of Participants", fontsize=30, labelpad=20)
plt.xlabel("VAC from RF Regressor", fontsize=30, labelpad=20)
plt.savefig(f'./Figures/hist_all_participants.png', bbox_inches="tight")
plt.show()
    
#%%