from pathlib import Path
print('Running' if __name__ == '__main__' else 'Importing', Path(__file__).resolve())
from datetime import timedelta
import datetime
from turtle import left
from sklearn.preprocessing import MinMaxScaler
from sklearn.preprocessing import RobustScaler
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Sequential 
import tensorflow as tf

from keras.layers import LSTM
from keras.layers import Dropout
from keras.models import Model
from keras.layers import Input, Dense
from keras.models import load_model
from keras.callbacks import EarlyStopping, ModelCheckpoint
import keras.backend as K

import numpy as np
import pandas as pd
from scipy import stats
import matplotlib.pyplot as plt
from data_handler import DataHandler

class Forecast:
    def __init__(self, dataHandler: DataHandler):
        self.batch_size = 3
        self.lookback = 10
        self.delay = 1
        self.predicted_column = 'longtime_mean'
        self.dataHandler = dataHandler
        self.scaler = RobustScaler()


    
    def train_model(self, begin_of_training = None, end_of_training = None):
        
        df, _ = self.dataHandler.get_data_for_training_model(left_time_interval=begin_of_training, right_time_interval=end_of_training, predicted_column=self.predicted_column)
        self.num_of_features = len(df.columns) - 1
        self.df_train_norm = df.copy()


        self.df_train_norm[df.columns] = self.scaler.fit_transform(df)

        self.train_gen = self.generator(dataframe = self.df_train_norm, 
                      target_name = self.predicted_column, 
                      lookback = self.lookback,
                      delay = self.delay,
                      min_index = 0,
                      max_index = int(df.shape[0]*0.8),
                      step = 1,
                      shuffle = True,
                      batch_size = self.batch_size)

        self.valid_gen = self.generator(dataframe = self.df_train_norm, 
                            target_name = self.predicted_column, 
                            lookback = self.lookback,
                            delay = self.delay,
                            min_index = int(df.shape[0]*0.8),
                            max_index = None,
                            step = 1,
                            shuffle = False,
                            batch_size = self.batch_size)
        

        
        self.val_steps = int((self.df_train_norm.shape[0]*0.1 - self.lookback) // self.batch_size)
        # This is how many steps to draw from `train_gen`
        # in order to see the whole train set:
        self.train_steps = int((self.df_train_norm.shape[0]*0.9 - self.lookback) // self.batch_size)


        
    def generator(self, dataframe, target_name, lookback, delay, min_index, max_index,shuffle=False, batch_size=128, step=6):
    
        data = dataframe.values
        data = data.astype(np.float32)
        target_indx = dataframe.columns.get_loc(target_name)
        
        if max_index is None:
            max_index = len(data) - delay - 1
        i = min_index + lookback
        while 1:
            if shuffle:
                rows = np.random.randint(
                    min_index + lookback, max_index, size=batch_size)
            else:
                if i + batch_size >= max_index:
                    i = min_index + lookback
                rows = np.arange(i, min(i + batch_size, max_index))
                i += len(rows)

            samples = np.zeros((len(rows),
                            lookback // step,
                            data.shape[-1]))
            targets = np.zeros((len(rows),))
            for j, row in enumerate(rows):
                indices = range(rows[j] - lookback, rows[j], step)
                samples[j] = data[indices]
                targets[j] = data[rows[j] + delay][target_indx]
            yield samples, targets
            
    def r2_keras(self, y_true, y_pred):
        """Coefficient of Determination 
        """
        SS_res =  K.sum(K.square( y_true - y_pred ))
        SS_tot = K.sum(K.square( y_true - K.mean(y_true) ) )
        return ( 1 - SS_res/(SS_tot + K.epsilon()) )
    
    def build_model(self):
        
        model = Sequential()
        model.add(tf.keras.Input(shape=(None, self.df_train_norm.shape[1])))
        model.add(tf.keras.layers.Conv1D(filters=6, kernel_size=5, activation='relu'))
        model.add(LSTM(6, return_sequences=True, activation='relu'))
        model.add(LSTM(6, return_sequences=False, activation='relu'))
        model.add(Dense(1))

        self.model = model
    
    def fit_model(self):
        callbacks = [EarlyStopping(monitor='loss', min_delta = 0, patience=10, verbose=1, mode='auto', restore_best_weights=True),
             ModelCheckpoint(filepath='lstm_model.h5', monitor='val_loss', save_best_only=True)]

        self.model.compile(loss='mae', optimizer='adam',metrics=[self.r2_keras])
        # history = model.fit(train_gen, epochs=50, batch_size=72, validation_data=valid_gen, verbose=2, shuffle=False, use_multiprocessing=True)

        history = self.model.fit(self.train_gen,
                                    steps_per_epoch=self.train_steps,
                                    epochs=70,
                                    shuffle=False,
                                    validation_data=self.valid_gen,
                                    validation_steps=self.val_steps,
                                    callbacks = callbacks)
        
    def get_forecast_next_steps(self, left_time_interval = None, right_time_interval = None, predicted_column = 'longtime_mean'):
        
        if (left_time_interval is None):
            left_time_interval = datetime.now() - timedelta(days=2)
        if (right_time_interval is None):
            right_time_interval = datetime.now()        
        
        df_all = self.dataHandler.get_data_for_prediction(left_time_interval=left_time_interval, right_time_interval=right_time_interval, predicted_column=self.predicted_column)

        forecast_future = pd.DataFrame()

        current_forecast_begin_date = left_time_interval
        current_forecast_end_date = left_time_interval + timedelta(minutes=30)
        for i in range(0, 48):
            print("-----")

            df_predict = pd.DataFrame({'datetime': pd.date_range(current_forecast_begin_date, current_forecast_end_date, freq='30min')})
            df_predict['longtime_mean'] = 0
            df_predict['weekday_sin'] = np.sin(2 * np.pi * df_predict['datetime'].dt.weekday / 7)
            df_predict['weekday_cos'] = np.cos(2 * np.pi * df_predict['datetime'].dt.weekday / 7)
            df_predict['hour_sin'] = np.sin(2 * np.pi * df_predict['datetime'].dt.hour / 24)
            df_predict['hour_cos'] = np.cos(2 * np.pi * df_predict['datetime'].dt.hour / 24)
            df_predict['minute_sin'] = np.sin(2 * np.pi * df_predict['datetime'].dt.minute / 60)
            df_predict['minute_cos'] = np.cos(2 * np.pi * df_predict['datetime'].dt.minute / 60)
            # delete column datetime
            df_predict = df_predict.drop(columns='datetime')
            len_df_predict = df_predict.shape[0]
            
            
            # concar df_all and df_predict
            
            df_all = pd.concat([df_all, df_predict], axis=0)
            df_all = df_all.reset_index(drop=True)
            
            df_predict_norm = df_all.copy()
            df_predict_norm[df_all.columns] = self.scaler.transform(df_all)
            # create predict df with values 

            predict_gen = self.generator(dataframe = df_predict_norm, 
                    target_name = self.predicted_column, 
                    lookback = self.lookback,
                    delay = self.delay,
                    min_index = 0,
                    max_index = None,
                    step = 1,
                    shuffle = False,
                    batch_size = df_predict.shape[0])

            (X, y_truth) = next(predict_gen)

            y_pred = self.model.predict(X)

            # np.expand_dims(y_truth,axis=1).shape
            y_pred_inv = np.concatenate((y_pred,np.zeros((y_pred.shape[0],self.num_of_features))),axis=1)
            y_pred_inv = self.scaler.inverse_transform(y_pred_inv)
            y_pred_inv = y_pred_inv[:,0]
            
            # set df_all last len_df_predict values to y_pred_inv
            df_all.iloc[-len_df_predict:, df_all.columns.get_loc(predicted_column)] = y_pred_inv
            # plt.plot(df_all.iloc[-len_df_predict:, df_all.columns.get_loc('longtime_mean')], color = 'green', label = 'Predicted data')
            df_all = df_all[len_df_predict:]
            forecast_future = pd.concat([forecast_future, df_all[-len_df_predict:]], axis=0)
            forecast_future = forecast_future.reset_index(drop=True)
            
            current_forecast_begin_date = current_forecast_begin_date + timedelta(hours=0.5)
            current_forecast_end_date = current_forecast_end_date + timedelta(hours=0.5)
        
        # create a dataframe with forecast and datetime as index
        return forecast_future
    


        
if __name__ == '__main__':
    dataHandler = DataHandler(
        "localhost",
        "smart_home_formankovi",
        "root",
        "root",
        "shelly1pm_34945475a969",
        "esphome_web_c771e8_tmp3"
    )    
    forecast = Forecast(dataHandler)
    forecast.train_model()
    forecast.build_model()
    forecast.fit_model()
    forecast.get_forecast_next_steps(30)
        
    