import pandas as pd
import xgboost
from sklearn.metrics import mean_squared_error, mean_absolute_percentage_error
import time

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


'''
Ajout de 5 valeurs anterieur
Et de 1 valeur posterieur : label de prediction
'''
def prepare_feat(df):
    feature = ['p', 'p_l1', 'p_l2', 'p_l3']
    label = 'p+1'
    nb_hist = 6
    for i in range(1,nb_hist):
        df = df.join(df.shift(periods=i)['p_l1'].to_frame(f'p_l1-{i}'))
        df = df.join(df.shift(periods=i)['p_l2'].to_frame(f'p_l2-{i}'))
        df = df.join(df.shift(periods=i)['p_l3'].to_frame(f'p_l3-{i}'))
        feature.extend((f'p_l1-{i}',f'p_l2-{i}',f'p_l3-{i}'))
    df = df.join(df.shift(periods=-1)['p'].to_frame(f'p+1'))
    df = df.dropna()
    return df, df[feature], df[label]


def predict(df, reg = None):
    if not reg:
        reg = xgboost.XGBRegressor()
        reg.load_model("model_xgboost.json")

    df, x, y = prepare_feat(df)

    for i in range(len(x)):
        tmp = x.iloc[i].to_frame().T
        predictions = reg.predict(tmp)
        error = round(mean_absolute_percentage_error([y.iloc[i]], predictions), 3)
        if error > 0.2:
            print(df.iloc[i].name, predictions, round(y.iloc[i],3), f'{bcolors.FAIL}{error}{bcolors.ENDC}')
            # TODO
            # send bad prediction and real value to cloud
        elif error > 0.1:
            print(df.iloc[i].name, predictions, round(y.iloc[i],3), f'{bcolors.WARNING}{error}{bcolors.ENDC}')
        else:
            print(df.iloc[i].name, predictions, round(y.iloc[i],3), f'{bcolors.OKGREEN}{error}{bcolors.ENDC}')
        
        time.sleep(0.1)
    


if __name__ == "__main__":

    df = pd.read_csv('159_2023-01-11_1d.csv')
    df['timestamp'] = pd.to_datetime(df['timestamp'])
    df = df.set_index('timestamp')
    print('predict')
    predict(df)