import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer
from lda import LDA

def learn_topics(texts, topicnum):

    # Get vocabulary and word counts.  Use the top 10,000 most frequent
    # lowercase unigrams with at least 3 alphabetical, non-numeric characters,
    # punctuation treated as separators.
    print("Vectorizing...")
    CVzer = CountVectorizer(max_features=10000,
                            lowercase=True)
    doc_vcnts = CVzer.fit_transform(texts)
    vocabulary = CVzer.get_feature_names()

    # Learn topics.  Refresh conrols print frequency.
    print("LDA")
    lda_model = LDA(topicnum, n_iter=4000, refresh=500) 
    doc_topic = lda_model.fit_transform(doc_vcnts)
    topic_word = lda_model.topic_word_

    return doc_topic, topic_word, vocabulary

print("Reading data...")
env = pd.read_csv('../Data/Environmental Discourse/env_processed.csv', index_col=0)
env = env[~env.text_processed.isna()]

doc_topic, topic_word, vocabulary = learn_topics(env.text_processed, 100)

print(doc_topic[0,:])

for i in range(100):
    env['topic_{}'.format(i)] = doc_topic[:, i]

env.to_csv('../Data/Environmental Discourse/env_lda.csv')