from typing import Optional, Tuple

import pandas as pd
import numpy as np
from scattertext.termscoring.CorpusBasedTermScorer import CorpusBasedTermScorer
from scipy.stats import rankdata

from scattertext.termranking.AbsoluteFrequencyRanker import AbsoluteFrequencyRanker
from scattertext.termscoring.ScaledFScore import ScaledFScorePresetsNeg1To1


class TermCategoryRanker(object):
    def __init__(
            self,
            scorer: type = ScaledFScorePresetsNeg1To1,
            term_ranker: type = AbsoluteFrequencyRanker,
            use_non_text_features: bool = False,
            target_category: Optional[str] = None
    ):
        self.scorer = scorer
        self.term_ranker = term_ranker
        self.use_non_text_features = use_non_text_features
        self.target_category = target_category

    def set_use_non_text_features(self, use_non_text_features: bool) -> 'TermCategoryRanker':
        self.use_non_text_features = use_non_text_features
        return self

    def get_rank_df(self, term_doc_matrix) -> pd.DataFrame:
        tdf, tdf_sum = self.__get_term_ranks_and_frequencies(term_doc_matrix)
        score_data = {}
        categories = term_doc_matrix.get_categories()
        my_scorer = self._initialize_scorer(term_doc_matrix)

        if self.target_category is not None:
            assert self.target_category in categories
            if issubclass(self.scorer, CorpusBasedTermScorer):
                score_data[str(self.target_category)] = my_scorer.set_categories(
                    category_name=self.target_category
                ).get_scores()
                score_data['Not ' + str(self.target_category)] = -1 * score_data[str(self.target_category)]
            else:
                focus = tdf[str(self.target_category)].values
                background = tdf_sum.values - focus
                score_data[str(self.target_category)] = my_scorer.get_scores(focus, background)
                score_data['Not ' + str(self.target_category)] = my_scorer.get_scores(background, focus)
        else:
            for category in categories:
                if issubclass(self.scorer, CorpusBasedTermScorer):
                    score_data[str(category)] = my_scorer.set_categories(
                        category_name=category
                    ).get_scores()
                else:
                    focus = tdf[str(category)].values
                    background = tdf_sum.values - focus
                    score_data[str(category)] = my_scorer.get_scores(focus, background)
        return pd.DataFrame(score_data, index=tdf.index).apply(lambda x: rankdata(x, 'dense'))

    def _initialize_scorer(self, term_doc_matrix):
        if issubclass(self.scorer, CorpusBasedTermScorer):
            my_scorer = self.scorer(term_doc_matrix)
            if self.use_non_text_features:
                return my_scorer.use_metadata()
            return my_scorer
        return self.scorer()

    def __get_term_ranks_and_frequencies(self, term_doc_matrix) -> Tuple[pd.DataFrame, pd.Series]:
        ranker = self.term_ranker(term_doc_matrix)
        if self.use_non_text_features:
            ranker = ranker.use_non_text_features()
        tdf = ranker.get_ranks('')
        tdf_sum = tdf.sum(axis=1)
        return tdf, tdf_sum

    def get_frequencies(self, term_doc_matrix) -> pd.Series:
        return self.__get_term_ranks_and_frequencies(term_doc_matrix)[1]

    def get_max_rank(self, term_doc_matrix) -> float:
        '''

        :param term_doc_matrix: TermDocMatrix
        :return: int
        '''
        rank_df = self.get_rank_df(term_doc_matrix)
        return rank_df.max().max()


class BaseAssociationCompactor(object):
    def __init__(self,
                 scorer=ScaledFScorePresetsNeg1To1,
                 term_ranker=AbsoluteFrequencyRanker,
                 use_non_text_features=False,
                 target_category: Optional[str] = None):
        self.scorer = TermCategoryRanker(scorer, term_ranker, use_non_text_features, target_category)

    def set_use_non_text_features(self, use_non_text_features: bool) -> 'BaseAssociationCompactor':
        self.scorer = self.scorer.set_use_non_text_features(use_non_text_features)
        return self

    def _prune_higher_ranked_terms(self, term_doc_matrix, rank_df, rank):
        terms_to_remove = rank_df.index[np.isnan(rank_df[rank_df <= rank])
        .apply(lambda x: all(x), axis=1)]
        return self._remove_terms(term_doc_matrix, terms_to_remove)

    def _remove_terms(self, term_doc_matrix, terms_to_remove):
        return term_doc_matrix.remove_terms(terms_to_remove, non_text=self.scorer.use_non_text_features)


class JSDCompactor(BaseAssociationCompactor):
    def __init__(self,
                 max_terms,
                 term_ranker=AbsoluteFrequencyRanker,
                 use_non_text_features=False,
                 target_category: Optional[str] = None):
        self.max_terms = max_terms
        BaseAssociationCompactor.__init__(
            self,
            term_ranker=term_ranker,
            use_non_text_features=use_non_text_features,
            target_category=target_category
        )

    def compact(self, term_doc_matrix, verbose=False):

        rank_df = self.scorer.get_rank_df(term_doc_matrix)
        p_df = rank_df / rank_df.sum(axis=0) + 0.001
        m = p_df.sum(axis=1)

        def lg(x): return np.log(x) / np.log(2)

        rank_df['Score'] = m * lg(1 / m) - (p_df * lg(1 / p_df)).sum(axis=1)
        terms_to_remove = rank_df.sort_values(
            by='Score', ascending=False
        ).iloc[self.max_terms:].index
        return term_doc_matrix.remove_terms(terms_to_remove, self.scorer.use_non_text_features)


class AssociationCompactor(BaseAssociationCompactor):
    def __init__(self,
                 max_terms,
                 scorer=ScaledFScorePresetsNeg1To1,
                 term_ranker=AbsoluteFrequencyRanker,
                 use_non_text_features=False,
                 include_n_most_frequent_terms=0,
                 target_category: Optional[str] = None):
        self.max_terms = max_terms
        self.include_n_most_frequent_terms = include_n_most_frequent_terms
        BaseAssociationCompactor.__init__(self, scorer, term_ranker, use_non_text_features, target_category)

    def compact(self, term_doc_matrix, verbose=False):
        '''
        Parameters
        ----------
        term_doc_matrix : TermDocMatrix
            Term document matrix object to compact
        Returns
        -------
        New term doc matrix
        '''

        rank_df = self.scorer.get_rank_df(term_doc_matrix)
        optimal_rank = self._find_optimal_rank(rank_df)

        terms_to_remove = rank_df.index[np.isnan(rank_df[rank_df <= optimal_rank])
        .apply(lambda x: all(x), axis=1)]

        if self.include_n_most_frequent_terms > 0:
            most_frequent_terms = self.scorer.get_frequencies(term_doc_matrix).sort_values(
                ascending=False
            ).index[:self.include_n_most_frequent_terms]
            terms_to_remove = set(terms_to_remove) - set(most_frequent_terms)

        compacted_term_doc_matrix = self._remove_terms(term_doc_matrix, terms_to_remove)

        if verbose:
            #print('max terms', self.max_terms,
            #      'optimal_rank', optimal_rank,
            #      'num_terms', compacted_term_doc_matrix.get_num_terms())
            pass
        return compacted_term_doc_matrix

    def _get_num_terms_at_rank(self, rank_i, rank_df):
        return sum(np.isnan(rank_df[rank_df <= rank_i]).apply(lambda x: not all(x), axis=1))

    def _find_optimal_rank(self, ranks_df):
        max_rank = ranks_df.max().max()
        min_rank = 1
        last_max_rank = None
        last_min_rank = None
        while max_rank - 1 > min_rank:
            if last_max_rank is not None:
                if last_min_rank == min_rank and last_max_rank == max_rank:
                    raise Exception("Error. Potential infinite loop detected.")
            last_max_rank = max_rank
            last_min_rank = min_rank
            cur_rank = int((max_rank - min_rank) / 2) + min_rank
            num_terms = self._get_num_terms_at_rank(cur_rank, ranks_df)
            if num_terms > self.max_terms:
                max_rank = cur_rank
            elif num_terms < self.max_terms:
                min_rank = cur_rank
            else:
                return cur_rank
        return min_rank


def find_optimal_rank(rank_df, num_terms):
    max_rank = rank_df.max().max()
    min_rank = 1
    last_max_rank = None
    last_min_rank = None
    while max_rank - 1 > min_rank:
        if last_max_rank is not None:
            if last_min_rank == min_rank and last_max_rank == max_rank:
                raise Exception("Error. Potential infinite loop detected.")
        last_max_rank = max_rank
        last_min_rank = min_rank
        cur_rank = int((max_rank - min_rank) / 2) + min_rank
        cur_num_terms = sum(np.isnan(rank_df[rank_df <= cur_rank]).apply(lambda x: not all(x), axis=1))
        if cur_num_terms > num_terms:
            max_rank = cur_rank
        elif cur_num_terms < num_terms:
            min_rank = cur_rank
        else:
            return cur_rank
    return min_rank


class AssociationCompactorByRank(BaseAssociationCompactor):
    def __init__(self,
                 rank,
                 scorer=ScaledFScorePresetsNeg1To1,
                 term_ranker=AbsoluteFrequencyRanker,
                 use_non_text_features=False,
                 target_category: Optional[str] = None):
        self.rank = rank
        BaseAssociationCompactor.__init__(self, scorer, term_ranker, use_non_text_features, target_category)

    def compact(self, term_doc_matrix):
        '''
        Parameters
        ----------
        term_doc_matrix : TermDocMatrix
            Term document matrix object to compact
        Returns
        -------
        TermDocMatrix


        '''
        rank_df = self.scorer.get_rank_df(term_doc_matrix)
        return self._prune_higher_ranked_terms(term_doc_matrix, rank_df, self.rank)
