

from sklearn.base import BaseEstimator
from collections import Counter
import pandas as pd
from numpy import sum, nan, isnan
from ut.util.uiter import window


class NextElementPredictor(BaseEstimator):

    def predict(self, seqs):
        preds = self.predict_proba(seqs)
        return [max(pred, key=lambda key: pred[key]) for pred in preds]

    def predict_proba(self, seqs):
        return list(map(self._predict_proba_conditioned_on_recent_subseq, seqs))

    def _predict_proba_conditioned_on_recent_subseq(self, recent_subseq):
        raise NotImplementedError("Need to implement this method")


class MarkovNextElementPred(NextElementPredictor):

    _list_of_attributes_to_display = ['markov_window', 'empty_element', 'keep_stats_in_memory']

    def __init__(self, markov_window=2, empty_element=-1, keep_stats_in_memory=True):
        self.markov_window = markov_window
        self.keep_stats_in_memory = keep_stats_in_memory
        self.empty_element = empty_element
        self._empty_element_padding = [empty_element] * (self.markov_window - 1)

    @property
    def total_tuple_count(self):
        """
        :return: Number of observed window tuples (sum of values in self.snip_tuples_counter_)
        """
        if self.total_tuple_count_ is not None:
            return self.total_tuple_count_
        else:
            total_tuple_count_ = sum(self.snip_tuples_counter_.values())
            if self.keep_stats_in_memory:
                self.total_tuple_count_ = total_tuple_count_
            return total_tuple_count_

    @property
    def pair_prob(self):
        """
        :return: Series of probabilities (unsmoothed count ratios) indexed by snip pairs
        """
        if self.pair_prob_ is not None:
            return self.pair_prob_
        else:
            pair_prob_ = pd.Series(self.snip_tuples_counter_) / self.total_tuple_count
            if self.keep_stats_in_memory:
                self.pair_probs_ = pair_prob_
            return pair_prob_

    @property
    def element_prob(self):
        """
        :return: Series of snips probabilities (unsmoothed count ratios)
        """
        if self.element_prob_ is not None:
            return self.element_prob_
        else:
            element_prob_ = (self.pair_prob * self.total_tuple_count)
            element_prob_ = element_prob_.groupby(level=0).sum()
            element_prob_ = element_prob_.drop(labels=self.empty_element)
            # element_prob_ = element_prob_.iloc[
            #     element_prob_.index.get_level_values(level=0) != self.empty_element]
            element_prob_ /= element_prob_.sum()
            if self.keep_stats_in_memory:
                self.element_prob_ = element_prob_
            return element_prob_

    @property
    def conditional_prob(self):
        """
        :return: Series of probabilities of last element (level) conditional on previous ones (including empty elements)
        """
        if self.conditional_prob_ is not None:
            return self.conditional_prob_
        else:
            conditional_prob_ = self._drop_empty_elements_of_sr(self.pair_prob, levels=[self.markov_window - 1])
            conditional_levels = list(range(self.markov_window - 1))
            conditional_prob_ = conditional_prob_.div(
                conditional_prob_.groupby(level=conditional_levels).sum(), level=0)  # TODO: Only works for two levels
            if self.keep_stats_in_memory:
                self.conditional_prob_ = conditional_prob_
            return conditional_prob_


    @property
    def initial_element_prob(self):
        """
        :return: Series of snips probabilities (unsmoothed count ratios)
        """
        if self.initial_element_prob_ is not None:
            return self.initial_element_prob_
        else:
            initial_element_prob_ = self.pair_prob.xs(self.empty_element, level=0, drop_level=True)
            initial_element_prob_ /= initial_element_prob_.sum()
            if self.keep_stats_in_memory:
                self.initial_element_prob_ = initial_element_prob_
            return initial_element_prob_

    def fit(self, snips_list):
        # reset anything previously learned
        self._initialize_params()
        return self.partial_fit(snips_list)

    def partial_fit(self, snips_list):
        if not set(['snip_tuples_counter_']).issubset(list(self.__dict__.keys())):
            self._initialize_params()
        for snips in snips_list:
            self._partial_fit_of_a_single_snips(snips)
        return self

    def _initialize_params(self):
        """
        Initializes model params (the snip_tuples_counter_, etc.)
        :return: None
        """
        self.snip_tuples_counter_ = Counter()
        self._reset_properties()

    def _reset_properties(self):
        """
        Resets some properties that depend on snip_tuples_counter_ to be computed (is used when the later changes)
        These will be recomputed when requested.
        :return: None
        """
        self.total_tuple_count_ = None
        self.pair_prob_ = None
        self.element_prob_ = None
        self.initial_element_prob_ = None
        self.conditional_prob_ = None

    def _partial_fit_of_a_single_snips(self, snips):
        self._reset_properties()
        self.snip_tuples_counter_.update(window(self._empty_element_padding + list(snips) + self._empty_element_padding,
                                                n=self.markov_window))


    def _drop_empty_elements_of_sr(self, sr, levels=None, renormalize=False):
        if levels is None:
            levels = list(range(self.markov_window))
        for level in levels:
            sr = sr.drop(labels=self.empty_element, level=level)
        if renormalize:
            sr /= sr.sum()
        return sr

    def _predict_proba_conditioned_on_recent_subseq(self, recent_subseq):
        pass

    def __repr__(self):
        d = {attr: getattr(self, attr) for attr in self._list_of_attributes_to_display if attr in self.__dict__}
        d['total_tuple_count'] = self.total_tuple_count
        return self.__class__.__name__ + '\n' + str(d)
