import numpy as np
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras import layers


class AE_MLP2:
    """
    Implementation of AE_MLP2
    
    Parameters
    ----------
    slidingwindow : int
        Subsequence length to analyze.
    epochs : int, (default=10)
        Number of epochs for the training phase

    Attributes
    ----------
    decision_scores_ : numpy array of shape (n_samples - subsequence_length,)
        The anomaly score.
        The higher, the more abnormal. Anomalies tend to have higher
        scores. This value is available once the detector is
        fitted.
    """  
    def __init__(self, slidingWindow = 100, epochs = 10, verbose=0):
        self.slidingWindow = slidingWindow
        self.epochs = epochs
        self.verbose = verbose
        self.model_name = 'AE_MLP2'

    def fit(self, X_clean, X_dirty, ratio = 0.15):
        """Fit detector.
        
        Parameters
        ----------
        X_clean : numpy array of shape (n_samples, )
            The input training samples.
        X_dirty : numpy array of shape (n_samples, )
            The input testing samples.
        ratio : flaot, ([0,1])
            The ratio for the train validation split
        
        Returns
        -------
        self : object
            Fitted estimator.
        """

        TIME_STEPS =  self.slidingWindow
        epochs = self.epochs
        

        X_train = self.create_dataset(X_clean,TIME_STEPS)
        X_test = self.create_dataset(X_dirty,TIME_STEPS)
        
        X_train = MinMaxScaler().fit_transform(X_train.T).T
        X_test = MinMaxScaler().fit_transform(X_test.T).T

        model = Sequential()
        model.add(layers.Dense(32,  activation='relu'))
        model.add(layers.BatchNormalization())
        model.add(layers.Dense(16, activation='relu'))
        model.add(layers.BatchNormalization())
        model.add(layers.Dense(32, activation='relu'))
        model.add(layers.Dense(TIME_STEPS, activation='relu'))
 
        model.compile(optimizer='adam', loss='mse')     
        
        history = model.fit(X_train, X_train,
                        epochs=epochs,
                        batch_size=64,
                        shuffle=False,
                        validation_split=0.15,verbose=self.verbose,
                        callbacks=[EarlyStopping(monitor="val_loss", verbose=self.verbose, patience=5, mode="min")])

        test_predict = model.predict(X_test)
        test_mae_loss = np.mean(np.abs(test_predict - X_test), axis=1)
        nor_test_mae_loss = MinMaxScaler().fit_transform(test_mae_loss.reshape(-1,1)).ravel()
        score = np.zeros(len(X_dirty))
        score[self.slidingWindow//2:self.slidingWindow//2+len(test_mae_loss)]=nor_test_mae_loss
        score[:self.slidingWindow//2]=nor_test_mae_loss[0]
        score[self.slidingWindow//2+len(test_mae_loss):]=nor_test_mae_loss[-1]
        
        self.decision_scores_ = score
        
        return self
        
    
    # Generated training sequences for use in the model.
    def create_dataset(self, X, time_steps):
        output = []
        for i in range(len(X) - time_steps + 1):
            output.append(X[i : (i + time_steps)])
        return np.stack(output)