from mylda.lda import LDA


class HLDA:
    """An implementation of hlda"""

    def __init__(self, K=10, iter_max=200):
        self.lda_model = LDA(K=K, iter_max=iter_max)

    def fit(self, dataframe):
        self.df = dataframe
        if "content" in self.df.columns:
            self.lda_model.train(content_list=list(dataframe["content"]))
        if "tokens" in self.df.columns:
            self.lda_model.train(token_list=list(dataframe["tokens"]))

    def nth_topic_part(self, n):
        df_ = self.df.drop("content", axis=1)
        df_["tokens"] = self.lda_model.n_th_topic_part(n)
        docNotEmpty = df_.apply(lambda r: len(r["tokens"]) > 0, axis=1)
        return df_[docNotEmpty]
