import argparse
import lm_dataformat as lmd
import gensim
import itertools
import pandas as pd
import seaborn as sns
import numpy as np
import json
import networkx as nx
import matplotlib.pyplot as plt
import spacy
import fasttext

from math import pow
from multiprocessing import Pool
from collections import defaultdict
from gensim.corpora.dictionary import Dictionary
from spacy.lang.en import English
from gensim.models import LdaModel, LdaMulticore
from sklearn.manifold import TSNE  
from best_download import download_file

component_list = {
    'StackExchange',
    'PubMed Central',
    'HackerNews',
    'OpenWebText2',
    'Enron Emails',
    'Ubuntu IRC',
    'BookCorpus',
    'USPTO',
    'Gutenberg (PG-19)',
    'Github',
    'CommonCrawl',
    'PubMed Abstracts',
    'FreeLaw',
    'Wikipedia (en)',
    'ArXiv',
    'DM Mathematics',
    'PhilPapers',
    'EuroParl',
    'Bibliotik',
    'OpenSubtitles',
    'YoutubeSubtitles',
    'NIH ExPorter',
}

component_options = component_list | { 'all' }

parser = argparse.ArgumentParser(description='Train topic models on a component.')

parser.add_argument('--input_path', required=True, help='Path to Pile split to use for dictionary (validation, ideally)')
parser.add_argument('--models_path', required=True, help='Path to saved LDA models')
parser.add_argument('--dictionary_file', required=True, help='File name to load the dictionary from')
parser.add_argument('--chunk_size', type=int, default=2**12,
                    help='Number of documents in a chunk')
parser.add_argument('--tokenizer_processes', type=int, default=1, help='Number of processes for input tokenization')
parser.add_argument('--lda_processes', type=int, default=1, help='Number of processes for LDA training')

args = parser.parse_args()

input_path = args.input_path
models_path = args.models_path
dictionary_file = args.dictionary_file
CHUNK_SIZE = args.chunk_size
TOK_PROCESSES = args.tokenizer_processes
LDA_PROCESSES = args.lda_processes

components = component_list

download_file('https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin', 'lid.176.bin', '7e69ec5451bc261cc7844e49e4792a85d7f09c06789ec800fc4a44aec362764e')
langdet = fasttext.load_model("lid.176.bin")
nlp = English()
tokenizer = nlp.Defaults.create_tokenizer(nlp)

def language(doc):
    details = langdet.predict(doc.replace('\n', ' '), k=1)

    return details[0][0].replace('__label__', '')

dictionary = Dictionary.load(dictionary_file)

rdr = lmd.Reader(input_path)
stream = rdr.stream_data(get_meta=True)

def baggify(item):
    text, meta = item
    component = meta['pile_set_name']
    if component in components:
        if language(text) == 'en':
            doc = tokenizer(text)
            tokens = [tok.lower_ for tok in doc if not tok.is_stop and tok.is_alpha]
            bow_or_none = dictionary.doc2bow(tokens)
        else:
            bow_or_none = None
    else:
        bow_or_none = None
    return (bow_or_none, component)

def raw_baggify(item):
    text, meta = item
    component = meta['pile_set_name']
    if component in components:
        if language(text) == 'en':
            doc = tokenizer(text)
            tokens_or_none = [tok.lower_ for tok in doc if not tok.is_stop and tok.is_alpha]
        else:
            tokens_or_none = None
    else:
        tokens_or_none = None
    return (tokens_or_none, component)

docs = defaultdict(list)
models = defaultdict()

print("Loading models...")
for component in components:
    models[component] = LdaMulticore.load(models_path + component + '.model.topic')

# Compute cross-corpus topic similarities    
diffs = defaultdict(defaultdict)

for (train, test) in itertools.product(*(components, components)):
    print(f'Computing topic diff for {train} on {test}...')
    diff, annotation = models[train].diff(models[test])
    diffs[train][test] = diff.tolist()

    fig, ax = plt.subplots(1,1)

    img = plt.imshow(diff, cmap='gray')

    ax.set_title("Topic Model Difference Matrix")
    ax.set_ylabel(train)
    ax.set_xlabel(test)
    
    plt.savefig(f'topic_diffs_{train}_{test}.png')
    
with open('topic_model_diffs.json', 'w') as fp:
    json.dump(diffs, fp)
    print("Saving topic modeling cross-corpus diffs...")

# Compute cross-corpus perplexities 
perplexities = defaultdict(defaultdict)

with Pool(TOK_PROCESSES) as p:
    print("Reading documents...")
    doc_iter = p.imap_unordered(baggify, stream, chunksize=CHUNK_SIZE)
    for (i, (bow, component)) in enumerate(doc_iter):
        if bow:
            docs[component].append(bow)    
            
for (train, test) in itertools.product(*(components, components)):
    print(f'Computing perplexities for model trained on {train} tested on {test}...')
    perplexities[train][test] = pow(2, -models[train].log_perplexity(docs[test]))
    
with open('topic_model_perplexities.json', 'w') as fp:
    json.dump(perplexities, fp)
    print("Saving topic modeling cross-corpus perplexities...")
    
df = pd.DataFrame.from_dict(perplexities, orient='index')
df = df.sort_index(axis=0)
df = df.sort_index(axis=1)
log_perplexities = df.apply(np.log)

print("Mean perplexities:", log_perplexities.mean(axis=1))

cm = sns.clustermap(log_perplexities.transpose(), dendrogram_ratio=0.33, metric='euclidean', xticklabels=True, yticklabels=True, col_cluster=False, cbar_pos=(0, 0, 0.05, 0.2), tree_kws={ 'linewidths': 4})
ax = cm.ax_heatmap
ax.set_xlabel("Training Corpus")
ax.set_ylabel("Testing Corpus")

fig = ax.get_figure()
fig.savefig('topic_model_perplexities.png')

# Show topics for each component
for component in components:
    print(component, models[component].show_topics(16, 10))
