

import collections
import numpy as np
import pandas as pd
import logging

class Losses(object):
    """
    Calculates prediction losses on test datasets achieved by the trained estimators. When the class is instantiated it creates a dictionary that stores the losses.

    Args:
        metric(`mlaut.analyze_results.scores` object): score function that will be used for the estimation. Must be `mlaut.analyze_results.scores` object.
        estimators(`array of mlaut estimators`): Array of estimators on which the results will be compared.
        exact_match(Boolean): If `True` when predictions for all estimators in the estimators array is not available no evaluation is performed on the remaining estimators. 
    """

    def __init__(self, metric):

        self._losses = collections.defaultdict(list)
        self._metric = metric
        self._losses_per_estimator = collections.defaultdict(list)
        self._losses_per_dataset_per_estimator = collections.defaultdict(list)
    def evaluate(self, predictions, true_labels, dataset_name, strategy_name):
        """
        Calculates the loss metrics on the test sets.

        Parameters
        ----------
        predictions: numpy array
            Predictions of trained estimators in the form 
        true_labels: numpy array
            true labels of test dataset.
        dataset_name: string
            Name of the dataset
        dataset_name: string
            Name of the strategy
        """

        #evaluates error per estimator
        loss = self._metric.calculate(true_labels, predictions)
        if strategy_name in self._losses_per_estimator:     
            self._losses_per_estimator[strategy_name].append(loss)
        else:
            self._losses_per_estimator[strategy_name] = [loss]

        #evaluate per dataset
        avg_score, std_score = self._metric.calculate_per_dataset(y_true=true_labels, 
                                                                  y_pred=predictions)

        
        self._losses_per_dataset_per_estimator[dataset_name].append([strategy_name, avg_score, std_score])
        
    def get_losses(self):
        """
        When the Losses class is instantiated a dictionary that holds all losses is created and appended every time the evaluate() method is run. This method returns this dictionary with the losses.

        Returns
        -------
            errors_per_estimator (dictionary), errors_per_dataset_per_estimator (dictionary), errors_per_dataset_per_estimator_df (pandas DataFrame): Returns dictionaries with the errors achieved by each estimator and errors achieved by each estimator on each of the datasets.  ``errors_per_dataset_per_estimator`` and ``errors_per_dataset_per_estimator_df`` return the same results but the first object is a dictionary and the second one a pandas DataFrame. ``errors_per_dataset_per_estimator`` and ``errors_per_dataset_per_estimator_df`` contain both the mean error and deviation.
        """ 
        return self._losses_per_estimator, self._losses_to_dataframe(self._losses_per_dataset_per_estimator)

    def _losses_to_dataframe(self, losses):
        """
        Reformats the output of the dictionary returned by the :func:`mlaut.analyze_results.losses.Losses.get_losses` to a pandas DataFrame. This method can only be applied to reformat the output produced by :func:`mlaut.analyze_results.Losses.evaluate_per_dataset`.

        Parameters
        ----------

        losses: dictionary returned by the :func:`mlaut.analyze_results.losses.Losses.get_losses` generated by :func:`mlaut.analyze_results.losses.Losses.evaluate_per_dataset`
        """

        df = pd.DataFrame(losses)
        #unpivot the data
        df = df.melt(var_name='dts', value_name='values')
        df['classifier'] = df.apply(lambda raw: raw.values[1][0], axis=1)
        df['loss'] = df.apply(lambda raw: raw.values[1][1], axis=1)
        df['std_error'] = df.apply(lambda raw: raw.values[1][2], axis=1)
        df = df.drop('values', axis=1)
        #create multilevel index dataframe
        dts = df['dts'].unique()
        estimators_list = df['classifier'].unique()
        score = df['loss'].values
        std = df['std_error'].values
        
        df = df.drop('dts', axis=1)
        df=df.drop('classifier', axis=1)
        
        df.index = pd.MultiIndex.from_product([dts, estimators_list])

        return df
