from lmi_proteins.indexes.BaseIndex import BaseIndex

import pandas as pd
import numpy as np

from typing import Dict, List, Tuple
from lmi_proteins.classifiers.enums import Classifiers
from lmi_proteins.classifiers.GMM import train_gmm_classifier
from lmi_proteins.classifiers.KMeans import train_kmeans_classifier
from lmi_proteins.utils import remove_key
from pickle import load, dump
import logging
import os
from lmi_proteins.utils import get_logger_config
from datetime import datetime
from collections import namedtuple
Data = namedtuple('Data', 'X y shapes')


class LMI(BaseIndex):
    """
    Learned Metric Index -- architecture containing series of interconnected models
    forming a tree over the dataset allowing fast similarity searching based on probabilities.

    This class is responsible for two basic operations:
    1. Building -- Training an series of interconnected models (classifiers) -- tree nodes,
                   keeping track of each training object's placing within the tree.
    2. Searching -- Navitaging through the structure by classifier's probability distribution,
                    collecting the answer.

    Attributes
    -------
    model_stack : List[List[Classifier]]
        Structure storing the trained classifiers -- the built LMI

    model_mapping : Dict[List[int], int]
        Mapping between the nodes and their classifier representation
        stored in `model_stack`.

    objects_in_buckets : Dict[Tuple[int], int]
        Mapping between identifier of buckets
        (= nodes in the lower-most level of the tree)
        and number of objects that each of these buckets contains.
    """
    def __init__(self, df: pd.DataFrame=None, labels_df=None, n_levels=None, load_from_file=False, model_path=""):
        
        def load_pickle(file):
            with open(file, "rb") as file:
                return load(file)
        
        if load_from_file:
            self.classifier = load_pickle(f'{model_path}/classifier.pkl') 
            pred_data = load_pickle(f'{model_path}/data.pkl') 
            self.data = Data(df, pred_data, (df.shape[0], pred_data.shape[0]))
            self.descriptor_values = load_pickle(f'{model_path}/descriptor_values.pkl') 
            self.labels = load_pickle(f'{model_path}/labels.pkl') 
            self.model_mapping = load_pickle(f'{model_path}/model_mapping.pkl') 
            self.model_stack = load_pickle(f'{model_path}/model_stack.pkl') 
            self.n_levels = load_pickle(f'{model_path}/n_levels.pkl') 
            self.objects_in_buckets = load_pickle(f'{model_path}/objects_in_buckets.pkl')
            self.pred_labels = load_pickle(f'{model_path}/pred_labels.pkl') 
            self.root_node_id = load_pickle(f'{model_path}/root_node_id.pkl') 
            self.supervised = load_pickle(f'{model_path}/supervised.pkl') 
            logging.basicConfig(level=logging.INFO, format=get_logger_config())
            self.LOG = logging.getLogger(__name__)
        else:
            super().__init__(
                df,
                labels_df=labels_df,
                n_levels=n_levels
            )

            self.model_stack = []
            self.model_mapping = {}
            self.objects_in_buckets = {}
            self.LOG.info(
                f'Created an empty LMI of {self.n_levels} levels to be built on a dataset of size: {df.values.shape}.'
            )
            self.classifier = None

    def save(self, model_name):
        def save_pickle(model_file_name, obj, objname):
            with open(f'saved_models/{model_file_name}/{objname}.pkl',"wb") as file_lmi:
                dump(obj, file_lmi)
        model_file_name = model_name
        os.makedirs(f"saved_models/{model_file_name}", exist_ok=True)
        
        save_pickle(model_file_name, self.classifier, "classifier")
        save_pickle(model_file_name, self.data.y, "data")
        save_pickle(model_file_name, self.descriptor_values, "descriptor_values")
        save_pickle(model_file_name, self.labels, "labels")
        save_pickle(model_file_name, self.model_mapping, "model_mapping")
        save_pickle(model_file_name, self.model_stack, "model_stack")
        save_pickle(model_file_name, self.n_levels, "n_levels")
        save_pickle(model_file_name, self.objects_in_buckets, "objects_in_buckets")
        save_pickle(model_file_name, self.pred_labels, "pred_labels")
        save_pickle(model_file_name, self.root_node_id, "root_node_id")
        save_pickle(model_file_name, self.supervised, "supervised")
    
    # --------------------------------------------------------------
    #                           BUILDING
    # --------------------------------------------------------------

    def train(self, model_config: Dict[str, Dict[str, str]], rebuild=False):  # noqa: C901
        """ Trains (builds) the whole LMI.

        The bulding starts from the root node (== 0-th level)
        and continues sequentially towards the leaf models.
        During the training, the following attributes are modified:
            - new models are appended to `self.model_stack`, each
              item is composed of another list and corresponds to one training level.
            - after each level is trained, the predictions are collected and
              saved in `self.data.y`, which are then used to split data for training
              the following level
            - `self.model_mapping` saves the data group identification so that we can map
              it to the corresponding model in `self.model_stack`

        After the training the LMI structure is saves in `self.model_stack`, information
        about occupied buckets is in `self.objects_in_buckets`.

        Parameters
        -------
        model_config: Dict[Dict[str, str]]
            Dictionary of models specification in the form:
                {{'level-0' : {'model': 'RF', 'n_estimators': 200 ...} ...}
            Loaded from `config/model-*.yml`

        rebuild: boolean [Optional]
            Decides whether to rebuild the saved (already trained) LMI structure if not empty
        """
        if not rebuild and len(self.model_stack) != 0:
            self.LOG.warn('Cannot overwrite contents of model_stack.'
                          'Set rebuild=True to re-train the LMI again.')
            return
        else:
            self.model_stack = []
            self.model_mapping = {}

        self.LOG.info(
            'Training model M.0 (root) on dataset'
            f"{self.data.X.shape} with {model_config['level-0']}."
        )
        self.classifier = model_config['level-0']

        predictions, model = self.train_node(
            model_config['level-0'],
            self.data.X.values,
            self.data.y[self.labels[0]].values,
            (())
        )
        self.data.y[f'{self.labels[0]}_pred'] = predictions
        self.model_stack.append(model)

        def get_model_name(model_name: any, level: int) -> str:
            """ Creates the model's name from int/tuple as a dot-separated string.

            Parameters
            -------
            model_name : int or Tuple
                A single integer or a sequence of integers (Tuple) identifying a node.
            level : int
                Current training level.
            Returns
            -------
            model_name: str
                Dot-separated string model name.
            """
            model_name_str = f'M.{level}.'
            model_name_str += '.'.join([str(x) for x in model_name])
            return model_name_str

        def get_trainable_data(
            current_group: pd.DataFrame
        ) -> Tuple[pd.DataFrame, pd.DataFrame]:
            """ Creates a separation of the current dataset to a part that can be
            further trained (has a label at the corresponding level) and those that cannot.

            Parameters:
                current_group (pd.DataFrameGroupby):
                    The whole group of data to be trained on.

            Returns:
                trainable, non_trainable group (Tuple[pd.GroupBy, pd.GroupBy])
            """
            if self.supervised:
                trainable = current_group.dropna()
                non_trainable = \
                    current_group.loc[~current_group.index.isin(current_group.dropna().index)]
            else:
                trainable = current_group
                non_trainable = pd.DataFrame([])

            return trainable, non_trainable

        for level, level_label in enumerate(self.labels[1:], start=1):
            self.LOG.info(f"Training level {level} with {model_config[f'level-{level}']}.")
            data_groups = self.create_data_groups(level)
            level_object_ids = []
            level_predictions = []
            level_model_stack = []
            model_mapping_level = []

            for group_name, group in data_groups:
                if level == 1:
                    group_name = tuple((group_name,))

                trainable, not_trainable = get_trainable_data(group[level_label])
                # If there are more non_trainable data, stop the branch to pre-prune
                # the index tree (applicable to M-index only)-- make all data non-traininable.
                if not_trainable.shape > trainable.shape:
                    not_trainable = group[level_label]
                    trainable = pd.Series()

                X = self.data.X.loc[trainable.index].values
                if self.supervised:
                    y = np.array(self.data.y[level_label].loc[trainable.index].values, dtype=np.uint16)
                else:
                    y = np.empty((1, 1))
                    y[:] = np.nan

                if X.shape[0] > 0:
                    self.LOG.debug(
                        f'    Training model {get_model_name(group_name, level)}'
                        f' on {X.shape} samples.'
                    )

                    predictions, model = self.train_node(
                        model_config[f'level-{level}'],
                        X,
                        y,
                        group_name
                    )
                    model_mapping_level.append(group_name)
                    level_model_stack.append(model)
                else:
                    predictions = []

                level_object_ids.extend(trainable.index.to_list() + not_trainable.index.to_list())

                not_trainable_pred_labels = [np.nan for _ in range(not_trainable.shape[0])]
                level_predictions.extend(list(predictions) + not_trainable_pred_labels)
                self.LOG.info(f'Training {group_name}')

            pred_df = pd.DataFrame(
                level_predictions,
                index=level_object_ids,
                columns=[f'{level_label}_pred']
            )
            self.data.y[f'{level_label}_pred'] = \
                pred_df.reindex(self.data.y.index)[f'{level_label}_pred'].astype(pd.UInt16Dtype())

            self.model_mapping[level] = dict.fromkeys(model_mapping_level)
            self.model_mapping[level].update((k, i) for i, k in enumerate(self.model_mapping[level]))
            self.model_stack.append(level_model_stack)

        self.objects_in_buckets = self.get_object_count_in_buckets(self.pred_labels)
        self.LOG.info('Finished training the LMI.')

    def train_node(
        self,
        level_config: Dict[str, Dict[str, str]],
        X: np.ndarray,
        y: np.ndarray,
        node_id: Tuple[int]
    ) -> Tuple[List, Classifiers]:
        """ Trains a single node (= one classifier).

        1. Prepares the input data
        2. Instantiates the model to use for training
        3. Runs the training function
        4. Returns the trained model with predictions

        If the desired model is multilabel NN, target labels are composed of 30-NN as opposed to 1-NN.

        Parameters
        -------
        level_config: Dict[str, Dict[str,str]]
            Dictionary of models specification
        X: np.ndarray
            Training data
        y: np.ndarray
            Training labels, empty array in case of unuspervised learning
        node_id: Tuple[int]
            Current node identifier

        Returns
        -------
        predictions: np.ndarray
            Array of shape (n_objects x n_descriptor_values)
        classifier: Classifier
            Trained classifier representing the node

        """
        assert X.shape[1] == self.descriptor_values

        if X.shape[0] == 1 or (self.supervised and np.unique(y).shape[0] <= 1 and np.unique(y)[0] != np.nan):
            self.LOG.debug(
                '        Only one data-point or one class to train with.'
                ' Using the DummyClassifier.'
            )
            classifier = Classifiers['Dummy'].value(node_id)
        else:
            params = remove_key(level_config, 'model')
            classifier = Classifiers[level_config['model']].value(node_id, **params)

        predictions = []
        if type(classifier).__name__ == 'LMINeuralNetwork':
            classifier.add_dense_layer(np.unique(y).shape[0], activation='softmax')
            classifier.compile(params)
            predictions = classifier.train(X, y, params['epochs'])

        elif type(classifier).__name__ == 'LMIGaussianMixtureModel':
            predictions = train_gmm_classifier(classifier, params, X, node_id)

        elif type(classifier).__name__ == 'LMIKMeans':
            predictions = train_kmeans_classifier(classifier, params, X, node_id)

        if type(classifier).__name__ in ['LMIRandomForest', 'LMILogisticRegression', 'LMIDummyClassifier']:
            predictions = classifier.fit(X, y)

        return predictions, classifier

    def create_data_groups(self, level: int) -> pd.DataFrame.groupby:
        """ Applies a groupby operation of the dataset based on
            collected probabilities from the previous level.

        Parameters
        -------
        level : int
            Target training level.

        Returns
        -------
        data_groups: pd.DataFrame.groupby
            GroupBy DataFrame based on the predictions.
        """
        group_condition = [f"{self.labels[j]}_pred" for j in range(level)]
        return self.data.y.groupby(group_condition)

    # --------------------------------------------------------------
    #                           SEARCHING
    # --------------------------------------------------------------

    def get_probability_distribution(
        self,
        node_id: Tuple[int],
        query: np.ndarray,
        level: int
    ) -> np.ndarray:
        """ Collects probabilities returned by a classifier (= node) for a given query.

        Parameters
        -------
        node_id : Tuple[int]
            Identifier of the node searched
        query : np.ndarray (1D)
            Searched object

        Returns
        -------
        predictions : np.ndarray
            Discrete prob. distribution over the node children (= predicted classes).
        """

        if level == 0:
            model = self.model_stack[level]
            return model.predict(query)
        elif level <= len(self.model_mapping):
            model_stack_index = self.model_mapping[level].get(node_id)
            if model_stack_index is not None:
                model = self.model_stack[level][model_stack_index]
                assert model.node_id == node_id, \
                    f"Searched node ID: {node_id} and model's node ID: {model.node_id} don't match." \
                    f"model_stack_index: {model_stack_index} | level: {level}."

                try:
                    return model.predict(query)
                except:
                    return np.array([[[model.node_id + (node, ) for node in model.classes_][0], 1.0]], dtype=object)
        else:
            return []

    def search_node(
        self,
        priority_queue: np.ndarray,
        query: np.ndarray
    ) -> Tuple[np.ndarray, Tuple[int]]:
        """ Pop the top node from priority queue,
            collects prob. distribution for its children,
            adds to the priority queue and sorts it.

        Parameters
        -------
        priority_queue : List[Tuple[int]]
            Priority queue - structure deciding the next node to visit.
        query : np.ndarray
            Searched object

        Returns
        -------
        priority_queue : List[Tuple[int]]
            Modified Priority queue
        node_id : Tuple[int]
            Identifier of the popped node
        """

        popped_node, priority_queue = self.pop_node(priority_queue)
        node_label = self.get_node_label(popped_node)
        popped_node[0] = node_label
        level = 0 if node_label == -1 else len(node_label)
        self.LOG.debug(f'-> {node_label}, p={popped_node[1]} | PQ: {priority_queue[:3]}')

        if level != self.n_levels:
            prob_distr = self.get_probability_distribution(node_label, query, level)
            if prob_distr is None or prob_distr == []:
                return priority_queue, popped_node
            #print(prob_distr, priority_queue)
            priority_queue = np.append(prob_distr, priority_queue, axis=0)
        else:
            return priority_queue, popped_node

        return self.sort_priority_queue(priority_queue), popped_node
