from gensim import corpora, models, similarities
import time
import pyLDAvis.gensim
import string
import nltk
import phrasemachine
import csv
from nltk.corpus import stopwords
import pandas as pd
import numpy as np
import json
import pickle
import sys

stops = set(stopwords.words("english"))
stops.add('said')
stops.add('')
stops.add('•')
stops.add('–')
stops.add('©')

def get_doc_topic(corpus, model, id_to_bow): #this gets the distribution per doc back.
    id_to_topic = {}
    for id in id_to_bow.keys():
        top_arr = model.__getitem__(id_to_bow[id], eps=0)
        id_to_topic[id] = top_arr
    return id_to_topic

def run_lda(id_to_text, n_topics):
    
    id_to_text = {i[0]:i[1].split(" ") for i in id_and_text}

    print("creating dictionaries...")
    dictionary = corpora.Dictionary(id_to_text.values())
    id_to_bow = {id:dictionary.doc2bow(text) for id, text in id_to_text.items()}
    corpus = [dictionary.doc2bow(text) for text in id_to_text.values()]

    print('starting LDA')
    start = time.time()
    lda = models.ldamodel.LdaModel(corpus,
                               num_topics=n_topics,
                               alpha='auto',
                               id2word=dictionary,
                               iterations=1000,
                               random_state=1)
    end = time.time()
    print('lda took: ', end-start, ' seconds')
    return lda, get_doc_topic(corpus, lda, id_to_bow), dictionary

def viz(lda, corpus, dictionary, html_fn):
    lda_visualization = pyLDAvis.gensim.prepare(lda, corpus, dictionary, sort_topics=False)
    pyLDAvis.save_html(lda_visualization, html_fn)

#main <----------------
k = int(sys.argv[2])
fn = sys.argv[1]
csv.field_size_limit(524288)
print("reading in text data...")
id_and_text = []
id_to_source = dict()
with open(fn, newline='') as data:

    reader = csv.reader(data)

    for line in reader:

        idn = line[0]
        text = line[1]
        id_to_source[idn] = line[2]

        text = text.replace('‘', '\'').replace('’', '\'').replace('“','"').replace('”','"').replace('—', '-').replace('…', '...').replace('\n', ' ')

        text = text.translate(str.maketrans(string.punctuation, ' ' * len(string.punctuation))).lower()

        word_list = text.split(" ")

        go_words = [word for word in [word for word in word_list if word not in stops]]

        text = ' '.join(go_words)

        id_and_text.append((idn, text))

lda, id_to_topic, dictionary = run_lda(id_and_text, k)


if (len(dictionary) != lda.get_topics().shape[1]):
    print("ERROR")

print(lda.get_topics().shape)
# lda.save(sys.argv[2])

article_count = 0
sources = dict()
for id in id_to_topic.keys():
    source = id_to_source[id]
    for v in id_to_topic[id]:
        if v[1] >= 0.6:
            if source not in sources.keys():
                sources[source] = np.zeros(k)
            sources[source][v[0]] += 1
            article_count += 1

d = dict()

d["mdsDat"] = dict()
d["tinfo"] = dict()
d["token.table"] = dict()

d["mdsDat"]["x"] = []
d["mdsDat"]["y"] = []
d["mdsDat"]["topics"] = []
d["mdsDat"]["cluster"] = []
d["mdsDat"]["Freq"] = np.zeros(k)

totals = np.zeros(k)
for i in range(k):
    for source in sources.keys():
        totals[i] += sources[source][i]

d["mdsDat"]["Freq"] = list((totals / sum(totals)) * 100)

term = []
freq = []
total = []
category = []
logprob = []
loglift = []

categories = ["Topic" + str(i + 1) for i in range(k)] + ["Default"]

source_order = sorted([(i, source) for i, source in enumerate(sources.keys())])

table = np.zeros((k, len(source_order)))

for i, source in source_order:
    table[:, i] = sources[source]

# print(article_count)
# print(sources)
# print(table)

for i in range(k):
    table[i, :] /= np.sum(table[i, :])

saliencies = dict()

for i, source in source_order:
    # P(t), marginal probability of source
    a = np.sum(sources[source]) / article_count

    distinctiveness = 0

    for l in range(k):

        # P(T | w), the likelihood that observed word w was generated by latent topic T
        b = sources[source][l] / totals[l]

        if b == 0:
            continue
        else:
            distinctiveness += b * np.log(b / a)

    saliency = a * distinctiveness

    saliencies[source] = saliency

saliency_order = sorted([(y, x) for x, y in saliencies.items()], reverse=True)

# print(saliency_order)

display_length = min(30, len(saliencies.keys()))

term = [0] * display_length
freq = [0] * display_length
total = [0] * display_length
category = [0] * display_length
logprob = [0] * display_length
loglift = [0] * display_length

for i in range(-1, k):
    for j, source in source_order:
        if i == -1:
            index = saliency_order.index((saliencies[source], source))

            if index < display_length:

                term[index] = source
                freq[index] = sum(sources[source])
                total[index] = sum(sources[source])
                category[index] = "Default"
                logprob[index] = abs(float(index) - display_length)
                loglift[index] = abs(float(index) - display_length)
        else:
            if table[i][j] > 0:
                term.append(source)
                freq.append(sources[source][i])
                total.append(sum(sources[source]))
                category.append(categories[i])
                phi = table[i][j]
                logprob.append(np.log(phi))
                loglift.append(np.log(phi / (np.sum(sources[source]) / article_count)))
            else:
                term.append(source)
                freq.append(sources[source][i])
                total.append(sum(sources[source]))
                category.append(categories[i])
                phi = table[i][j]
                logprob.append(-9999)
                loglift.append(-9999)

rel_topic = []
rel_freq = []
rel_term = []

for _, source in source_order:

    s = np.sum(sources[source])

    for i in range(k):
        if sources[source][i] > 0:
            rel_topic.append(i + 1)
            rel_freq.append(sources[source][i] / s)
            rel_term.append(source)

d["tinfo"]["Term"] = term
d["tinfo"]["Freq"] = freq
d["tinfo"]["Total"] = total
d["tinfo"]["Category"] = category
d["tinfo"]["logprob"] = logprob
d["tinfo"]["loglift"] = loglift
d["token.table"]["Topic"] = rel_topic
d["token.table"]["Freq"] = rel_freq
d["token.table"]["Term"] = rel_term
d["R"] = 1

raw = json.dumps(d)
part2 = raw[raw.index("\"Freq\":") + 8:raw.index(", \"R\":")]

old = None

with open(sys.argv[3] + '.html', 'r') as file:
    old = file.read()

part1 = old[:old.index("\"Freq\":") + 8]
part3 = old[old.index(", \"R\":"):]

new = part1 + part2 + part3

with open(sys.argv[3] + '_sources' + '.html', 'w') as file:
    file.write(new)

