from pathlib import Path
print('Running' if __name__ == '__main__' else 'Importing', Path(__file__).resolve())
from datetime import timedelta
import datetime
from sklearn.preprocessing import RobustScaler
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
from smartboiler.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)
        
    