import logging
from lmi_proteins.utils import get_logger_config

import numpy as np
import pandas as pd
from typing import List, Tuple, Any
import time
from collections import namedtuple
Data = namedtuple('Data', 'X y shapes')


class BaseIndex(object):
    """
    Class representing BaseIndex and main operation on it - training and searching.
    Sets up core attributes of an index - index shape, dataset and labels.

    Attributes
    ----------
    data : Data
        The dataset to work with in building the index and searching within it.
        For Supervised learning, `df` as well as `labels_df` should be provided.
        For Unsupervised learning `labels_df` is provided,
            a new one is created with null values and filled with predictions lates in index building.
        The Data attribute is then composed of a tuple (`df`, `labels_df`, original DataFrame shapes)
    n_levels : str
        Specified number of levels in LMI.
        For Supervised learning the value is derived from `labels_df`
        For Unsupervised learning, the value is determined from config's `LMI.n_levels`.
    descriptor_values : int
        The length of the descriptor vectors (`df.shape[1]`).
    labels : List[str]
        The labels of the individual levels ('L1'...'Ln' by default).

    """
    def __init__(self, df: pd.DataFrame, labels_df=None, n_levels=None):
        """
        Parameters
        ----------
        lmi_config : Dict
            Loaded configuration from `config/config.yml`
        df : pd.DataFrame
            Dataset to build the index on (i.e., object descriptors without labels)
        labels_df : pd.DataFrame
            [Optional] Dataset labels. Provided only in the case of Supervised LMI
        """
        if labels_df is None:
            labels_df = pd.DataFrame(
                columns=[f'L{i}' for i in range(1, n_levels+1)],
                index=df.index
            )
            self.supervised = False
        else:
            self.supervised = True
        self.n_levels = labels_df.shape[1]
        self.descriptor_values = df.shape[1]
        self.labels = sorted(labels_df.columns.to_list())
        self.pred_labels = [f'{label}_pred' for label in self.labels]

        assert df.shape[0] == labels_df.shape[0], 'Data and labels need to have the same number of rows.'
        labels_df = labels_df.reindex(index=df.index)
        assert all(df.index == labels_df.index), 'Data and labels need to have the same index (i.e., object_id).'
        self.data = Data(df, labels_df, (df.shape[0], labels_df.shape[0]))
        self.root_node_id = -1

        logging.basicConfig(level=logging.INFO, format=get_logger_config())
        self.LOG = logging.getLogger(__name__)

    def get_object_count_in_buckets(self, labels):
        """ Creates dictionary with object counts in individual buckets.

        Returns
        -------
        count: Dict[Tuple[int], int]
            Mapping between identifier of buckets
            and number of objects that each of these buckets contains.
        """
        nan_placeholder = np.iinfo(np.int16).max
        count = self.data.y.fillna(nan_placeholder).groupby(labels).size().reset_index(name='n')
        count['bucket'] = list(count[labels].itertuples(
            index=False, name=None
            )
        )
        count['bucket'] = count['bucket'].apply(
            lambda x: tuple([label for label in x if label != nan_placeholder])
        )
        return dict(zip(count['bucket'].values, count['n'].values))

    def sort_priority_queue(self, priority_queue, ascending=False):
        """ Sorts priority queue in descending order, unless specified otherwise.

        Parameters:
            priority_queue (np.ndarray):
                Priority queue - structure governing the search
            ascending (bool, optional):
                Dictates ascending sorting. Defaults to False.

        Returns:
            priority_queue (np.ndarray):
                Sorted Priority queue
        """
        sorted_arr = priority_queue[priority_queue[:, 1].argsort()[::-1]].astype(object)
        if ascending:
            return np.flip(sorted_arr, axis=0)
        else:
            return sorted_arr

    def get_object(self, object: Any, reshape=True) -> np.ndarray:
        """ Gets an object (one row) from the dataset.

        Parameters:
            object_id (int):
                Object's identifier
            reshape (bool, optional):
                Dictates whether the object should be reshaped
                from (n_descriptors, ) shape. Defaults to True.

        Returns:
            np.ndarray:
                Retrieved object.
        """
        if type(object) == np.ndarray:
            return object.reshape(1, -1)
        elif type(object) == pd.core.series.Series:
            return object.values.reshape(1, -1)
        else:
            try:
                object = self.data.X.loc[object, ].values
                return object.reshape(1, -1)
            except KeyError:
                raise Exception(f'Object with id {object} does not exist in the trained dataset.')

    def pop_node(
        self,
        priority_queue: np.array
    ) -> Tuple[np.ndarray, np.ndarray]:
        """ Pops node from the Priority queue

        Parameters:
            priority_queue (np.array):
                Priority queue - structure governing the search

        Returns:
            Tuple[np.ndarray, np.ndarray]:
                popped node in the form (node_label, probability), priority queue
        """
        popped_node = priority_queue[0]
        priority_queue = np.delete(priority_queue, 0, axis=0)
        return popped_node, priority_queue

    def get_node_label(self, popped_node: np.ndarray) -> Tuple[int]:
        """ Unifies a give `popped_node` label format as tuple of ints.

        Parameters:
            popped_node (np.ndarray):
                A popped node in the form (node_label, probability)

        Returns:
            Tuple[int]:
                Node label
        """
        node_label = popped_node[0]
        if isinstance(node_label, float) and node_label != -1:
            node_label = tuple((int(node_label),))
        return node_label

    def search(
        self,
        query_id: int,
        stop_conditions: List[int]
    ) -> Tuple[List[np.ndarray], List[float], List[int]]:
        """ Performs search on a build index with a given query until all of the
        (potentially) many stop conditions are met.

        Starts with 'popping' the root node and putting all of its children
        to a Priority queue (PQ) - sorted structure with probabilities
        indicating similarity to the query.
        Then continues popping the top-most nodes and putting their children
        into PQ, then re-sorting (`self.search_node`).

        Parameters:
            query_id (int): Query to be searched
            stop_conditions (List[int]):
                List of the number of objects in visited buckets.
                Used as a threshold to stop searching at this point.

        Returns:
            Search history, time checkpoints, number of visited objects
        """

        def initialize_priority_queue() -> np.ndarray:
            """
            Create an initial entry for the root model for the priority queue.

            Returns
            -------
            initial priority queue: np.array
            """
            return np.array((self.root_node_id, 1.0)).reshape(1, -1)

        start = time.time()

        priority_queue = initialize_priority_queue()
        objects_in_buckets_visited = 0
        current_ckpt = 0
        popped_nodes = []
        search_history = []
        time_checkpoints = []
        visited_objects = []
        query = self.get_object(query_id)

        while len(priority_queue) != 0:

            priority_queue, popped_node = \
                self.search_node(priority_queue, query)

            popped_nodes.append(popped_node[0])

            if popped_node[0] != -1 and popped_node[0] in self.objects_in_buckets:

                objects_in_buckets_visited += \
                    self.objects_in_buckets[popped_node[0]]

                while objects_in_buckets_visited >= stop_conditions[current_ckpt]:

                    bucket_nodes = [p for p in popped_nodes if p in self.objects_in_buckets]
                    search_history.append(bucket_nodes)
                    time_checkpoints.append(time.time() - start)
                    visited_objects.append(objects_in_buckets_visited)

                    self.LOG.info(
                        f'Finished searching for query {query_id}, '
                        f'stop condition={stop_conditions[current_ckpt]}, '
                        f'visited {objects_in_buckets_visited} objects in total, '
                        f'took {round(time_checkpoints[-1], 2)}s.'
                    )
                    if current_ckpt == len(stop_conditions) - 1:
                        return search_history, time_checkpoints, visited_objects
                    else:
                        current_ckpt += 1
        return priority_queue
