import torch.nn as nn
import numpy as np
from tqdm import tqdm
import time

from functools import reduce

from sklearn import metrics
from sklearn.base import BaseEstimator, OutlierMixin
from sklearn.pipeline import Pipeline, FeatureUnion

from mtsa.features.mel import Array2Mfcc
from mtsa.models.GANF_components.ganfBaseModel import GANFBaseModel
from mtsa.utils import Wav2Array
from mtsa.features.stats import FEATURES


"""
    GRAPH-AUGMENTED NORMALIZING FLOWS

    This class implements the model described in the paper:
        GRAPH-AUGMENTED NORMALIZING FLOWS FOR ANOMALY DETECTION OF MULTIPLE TIME SERIES

    Key concepts and techniques from the paper have been adapted and implemented here.
    Any deviations or modifications from the original methodology are noted in the implementation details.

    Reference:
    [1] X. Wu, Y. Park, S. Tirthapura, "GRAPH-AUGMENTED NORMALIZING FLOWS FOR ANOMALY DETECTION OF MULTIPLE TIME SERIES,"
    AAAI, vol. 36, no. 8, 2022.  DOI: 10.1609/aaai.v36i8.20829
"""

class GANF(nn.Module, BaseEstimator, OutlierMixin):
    def __init__(
        self,
        sampling_rate=None,
        random_state=None,
        isForWaveData=False,
        batch_size=None,
        learning_rate=None,
        mono=True,
        use_array2mfcc=False,
        device=0,
        use_featureUnion= False
    ) -> None:
        super().__init__()
        self.sampling_rate = sampling_rate
        self.random_state = random_state
        self.isForWaveData = isForWaveData
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.mono = mono
        self.use_featureUnion = use_featureUnion
        self.final_model = GANFBaseModel(device=device)
        self.use_array2mfcc = use_array2mfcc
        self.model = self._build_model()
        self.last_fit_time = 0

    @property
    def name(self):
        return "GANF " + "+".join([f[0] for f in self.features])

    def fit(
        self,
        X,
        y=None,
        batch_size=None,
        epochs=None,
        max_iteraction=None,
        learning_rate=None,
        debug_dataframe=None,
    ):
        mono = self.mono and not self.use_array2mfcc
        if batch_size is None:
            batch_size = self.batch_size
        if learning_rate is None:
            learning_rate = self.learning_rate
        
        start = time.perf_counter()    
        self.model.fit(
            X,
            y,
            final_model__batch_size=batch_size,
            final_model__epochs=epochs,
            final_model__max_iteraction=max_iteraction,
            final_model__learning_rate=learning_rate,
            final_model__debug_dataframe=debug_dataframe,
            final_model__isWaveData=self.isForWaveData,
            final_model__mono=mono,
        )
        end = time.perf_counter()
        
        self.last_fit_time = end - start
        return self.last_fit_time #seconds

    def transform(self, X, y=None):
        l = list()
        l.append(X)
        l.extend(self.model.steps[:-1])
        Xt = reduce(lambda x, y: y[1].transform(x), l)
        return Xt

    def predict(self, X):
        return self.model.predict(X)

    def score_samples(self, X):
        return np.array(
            list(
                map(
                    self.model.score_samples, 
                    [[x] for x in X])
                )
            )

    def score(self, X, Y):
        X = self.score_samples
        fpr, tpr, thresholds = metrics.roc_curve(Y, X)
        auc = metrics.auc(fpr, tpr)
        return auc

    def get_adjacent_matrix(self):
        return self.final_model.get_adjacent_matrix()

    def _build_model(self):
        array2mfcc = Array2Mfcc(sampling_rate=self.sampling_rate)
        wav2array = Wav2Array(sampling_rate=self.sampling_rate, mono=self.mono)
        features = FeatureUnion(FEATURES)
        # mono = (self.mono and not self.use_array2mfcc)
        
        if self.use_featureUnion and self.use_array2mfcc and self.isForWaveData:
            model = Pipeline(
                steps=[
                    ("wav2array", wav2array),
                    ("array2mfcc", array2mfcc),
                    ("features", features),
                    ("final_model", self.final_model),
                ]
            )
        elif self.use_array2mfcc and self.isForWaveData:
            model = Pipeline(
                steps=[
                    ("wav2array", wav2array),
                    ("array2mfcc", array2mfcc),
                    ("final_model", self.final_model),
                ]
            )
        elif self.isForWaveData:
            model = Pipeline(
                steps=[
                    ("wav2array", wav2array),
                    ("final_model", self.final_model),
                ]
            )
        else:
            model = Pipeline(
                steps=[
                    ("final_model", self.final_model),
                ]
            )

        return model

    def attach_observer(self, observer):
        if (self.final_model is not None):
            self.final_model.attach_observer(observer)
            
    def set_adjacent_matrix(self, adjacent_matrix):
        if (self.final_model is not None):
            self.final_model.set_adjacent_matrix(adjacent_matrix)
            
    def get_random_adjacent_matrix(self):
         return self.final_model.get_random_adjacent_matrix()
     
    def get_initial_adjacent_matrix(self):
        return self.final_model.get_initial_adjacent_matrix()
    
