from abc import ABCMeta, abstractmethod
import logging

import pandas as pd
from six import with_metaclass

# doctest-only imports
from topik.readers import read_input
from topik.tests import test_data_path
from topik.intermediaries.persistence import Persistor

registered_models = {}

def register_model(cls):
    global registered_models
    if cls.__name__ not in registered_models:
        registered_models[cls.__name__] = cls
    return cls


class TopicModelBase(with_metaclass(ABCMeta)):
    corpus = None

    @abstractmethod
    def get_top_words(self, topn):
        """Method should collect top n words per topic, translate indices/ids to words.

        Return a list of lists of tuples:
        - outer list: topics
        - inner lists: length topn collection of (weight, word) tuples
        """
        raise NotImplementedError

    @abstractmethod
    def save(self, filename, saved_data):
        self.persistor.store_model(self.get_model_name_with_parameters(),
                                   {"class": self.__class__.__name__,
                                    "saved_data": saved_data})
        self.corpus.save(filename)

    @abstractmethod
    def get_model_name_with_parameters(self):
        raise NotImplementedError

    @abstractmethod
    def _get_term_data(self):
        raise NotImplementedError

    @abstractmethod
    def _get_vocab(self):
        raise NotImplementedError

    @abstractmethod
    def _get_term_frequency(self):
        raise NotImplementedError

    @abstractmethod
    def _get_topic_term_dists(self):
        raise NotImplementedError

    @abstractmethod
    def _get_doc_data(self):
        raise NotImplementedError

    @abstractmethod
    def _get_doc_topic_dists(self):
        raise NotImplementedError

    @abstractmethod
    def _get_doc_lengths(self):
        raise NotImplementedError

    def to_py_lda_vis(self):
        doc_data_df = self._get_doc_data()
        term_data_df = self._get_term_data()

        model_lda_vis_data = {  'vocab': term_data_df['token'],
                                'term_frequency': term_data_df['doc_count'],
                                'topic_term_dists': term_data_df.iloc[:,2:].T,
                                'doc_topic_dists': doc_data_df.iloc[:,:-1],
                                'doc_lengths': doc_data_df['doc_length']}
        return model_lda_vis_data

    def termite_data(self, topn_words=15):
        """Generate the pandas dataframe input for the termite plot.

        Parameters
        ----------
        topn_words: int
            number of words to include from each topic

        >>> raw_data = read_input('{}/test_data_json_stream.json'.format(test_data_path), "abstract")
        >>> processed_data = raw_data.tokenize()  # tokenize returns a DigestedDocumentCollection
        >>> model = registered_models["LDA"](processed_data, ntopics=3)
        >>> model.termite_data(15)
        
        """
        from itertools import chain
        return pd.DataFrame(list(chain.from_iterable([{"topic": topic_id, "weight": weight, "word": word}
                                                      for (weight, word) in topic]
                                                     for topic_id, topic in enumerate(self.get_top_words(topn_words)))))

    @property
    def persistor(self):
        return self.corpus.persistor


def load_model(filename, model_name):
    """Loads a JSON file containing instructions on how to load model data.

    Returns a TopicModelBase-derived object."""
    p = Persistor(filename)
    if model_name in p.list_available_models():
        data_dict = p.get_model_details(model_name)
        model = registered_models[data_dict['class']](**data_dict["saved_data"])
    else:
        raise NameError("Model name {} has not yet been created.".format(model_name))
    return model

@@ -2828,111 +2828,935 @@
 %3E%3E%3E 
-model = registered_models%5B%22LDA%22%5D(processed_data, ntopics=3)%0A        %3E%3E%3E model.termite_data(15)%0A        
+# must set seed so that we get same topics each run%0A        %3E%3E%3E import random%0A        %3E%3E%3E import numpy%0A        %3E%3E%3E random.seed(42)%0A        %3E%3E%3E numpy.random.seed(42)%0A        %3E%3E%3E model = registered_models%5B%22LDA%22%5D(processed_data, ntopics=3)%0A        %3E%3E%3E model.termite_data(5)%0A            topic    weight         word%0A        0       0  0.005337           nm%0A        1       0  0.005193         high%0A        2       0  0.004622        films%0A        3       0  0.004457       matrix%0A        4       0  0.004194     electron%0A        5       1  0.005109   properties%0A        6       1  0.004654         size%0A        7       1  0.004539  temperature%0A        8       1  0.004499           nm%0A        9       1  0.004248   mechanical%0A        10      2  0.007994         high%0A        11      2  0.006458           nm%0A        12      2  0.005717         size%0A        13      2  0.005399    materials%0A        14      2  0.004734        phase%0A
 %0A   

fix termite_data doctest