# --------------------------------------------------------------------------------------
# File: "ModelOutput.py"
# Dir: "src/sandsldahelper/"
# Created: 2021-08-17
# --------------------------------------------------------------------------------------

"""
this file is meant to create output files for the topic model created from the data provided
"""

from pathlib import Path
from tqdm import tqdm
import SandsPythonFunctions.ParquetFunctions as pf
import msgpack
import numpy as np
import pandas as pd
import pretty_errors
import tomotopy as tp


def create_pyLDAVis(mdl, pyvs_path):
    import pyLDAvis
    import numpy as np

    if not pyvs_path.exists():
        topic_term_dists = np.stack([mdl.get_topic_word_dist(k) for k in range(mdl.k)])
        doc_topic_dists = np.stack([doc.get_topic_dist() for doc in mdl.docs])
        doc_topic_dists /= doc_topic_dists.sum(axis=1, keepdims=True)
        doc_lengths = np.array([len(doc.words) for doc in mdl.docs])
        vocab = list(mdl.used_vocabs)
        term_frequency = mdl.used_vocab_freq
        vis = pyLDAvis.prepare(
            topic_term_dists,
            doc_topic_dists,
            doc_lengths,
            vocab,
            term_frequency,
            start_index=0,  # tomotopy starts topic ids with 0, pyLDAvis with 1
            sort_topics=False,  # IMPORTANT: otherwise topic_ids between pyLDAvis and tomotopy will not match
        )
        pyLDAvis.save_html(vis, open(str(pyvs_path.resolve()), "w"))


def create_topic_summery_text_page(mdl, topt_path):
    if not topt_path.exists():
        extractor = tp.label.PMIExtractor(
            min_cf=10, min_df=5, max_len=5, max_cand=10000, normalized=True
        )
        cands = extractor.extract(mdl)
        labeler = tp.label.FoRelevance(mdl, cands, min_df=5, smoothing=1e-2, mu=0.25)
        output_string = ""
        for k in range(mdl.k):
            output_string = f"{output_string}== Topic #{k} ==\n"
            output_string = f"{output_string}\nLabels: {', '.join(label for label, score in labeler.get_topic_labels(k, top_n=5))}\n"
            for word, prob in mdl.get_topic_words(k, top_n=30):
                if len(word) < 4:
                    output_string = f"{output_string}\n{word}\t\t\t{prob}"
                elif len(word) > 7:
                    output_string = f"{output_string}\n{word}\t{prob}"
                else:
                    output_string = f"{output_string}\n{word}\t\t{prob}"
            output_string = f"{output_string}\n\n"
        topt_path.write_text(output_string)


def extract_dominant_topic(mdl, ltxp_path, dtdx_path, corp_path, dtdp_path, dtcx_path):
    if not (dtdp_path.exists() and dtdx_path.exists() and dtcx_path.exists()):
        dta = pf.read_parquet_by_path(ltxp_path)
        # print(f"len dta: {len(dta)}; len mdl.docs: {len(mdl.docs)}")  # TESTCODE:
        if len(mdl.docs) == len(dta):
            corpus = tp.utils.Corpus.load(str(corp_path))
            inferred_corpus, ll = mdl.infer(corpus)
            dom_dist = []
            for doc in tqdm(inferred_corpus, desc="creating individual post dominant topic output"):
                dom_dist.append(doc.get_topic_dist().tolist())
            # print(f"type(dom_dist[0]: {type(dom_dist[0])}")  # TESTCODE:
            dta["dom_dist"] = dom_dist
        else:
            lemmatized_texts = dta["lemmatized_texts"].tolist()
            dom_dist = []
            for lemmatized_text in tqdm(lemmatized_texts, desc="Processing dominant topic by post"):
                doc = lemmatized_text.split()
                try:
                    doc = mdl.make_doc(doc)
                    topic_dist, ll = mdl.infer(doc)
                    dom_dist.append(topic_dist.tolist())
                except:
                    dom_dist.append(np.nan)
            dta["dom_dist"] = dom_dist
            dta = dta.dropna()
            dom_dist = dta["dom_dist"].tolist()
        dta["dom_topic"] = [(dist.index(max(dist)) + 1) for dist in dom_dist]
        dta["dom_topic_score"] = [max(dist) for dist in dom_dist]
        dta = dta[["id", "text", "dom_dist", "dom_topic", "dom_topic_score"]]
        pf.save_dataframe_as_parquet(dta, dtdp_path)
        print(dta.columns)
        dta.to_excel(dtdx_path)
        dta = dta.groupby("dom_topic")["id"].count()
        dta.to_excel(dtcx_path)


def create_dominant_topic_dta(domt_path, dta, text_dom_topics, percents_dom_topic):
    dta["dominant_topic"] = text_dom_topics
    dta["percent_dominant_topic"] = percents_dom_topic
    pf.save_dataframe_as_parquet(dta[["id", "dominant_topic", "percent_dominant_topic"]], domt_path)
    dta["percent_dominant_topic_string"] = [
        str(round(per, 3)).ljust(5, "0") for per in percents_dom_topic
    ]
    dta["dominant_topic_string"] = [str(dom).zfill(3) for dom in text_dom_topics]
    return dta


def create_dominant_topic_text_file(dta, dtxt_path, tpex_path, tgex_path):
    dta["text_no_newline"] = dta["text"].str.replace("\n", " |NEWLINE_CHAR| ")
    dta.to_excel(tpex_path, sheet_name="dominant_topics")
    dta = dta.sort_values(
        ["dominant_topic_string", "percent_dominant_topic_string"], ascending=[True, False]
    )
    dta["output_string"] = (
        dta["dominant_topic_string"]
        + " - "
        + dta["percent_dominant_topic_string"]
        + " - "
        + dta["id"]
        + " - "
        + dta["text_no_newline"]
    )
    output_string = dta.output_string.str.cat(sep="\n\n")
    dtxt_path.write_text(output_string)


def create_top_sample_dominant_topic_text_file(dta, dtst_path):
    dta["text_no_newline"] = dta["text"].str.replace("\n", " |NEWLINE_CHAR| ")
    dta = dta.sort_values(
        ["dominant_topic_string", "percent_dominant_topic_string"], ascending=[True, False]
    )
    dta = dta.groupby("dominant_topic_string").head(4).reset_index(drop=True)
    dta["output_string"] = (
        dta["dominant_topic_string"]
        + " - "
        + dta["percent_dominant_topic_string"]
        + " - "
        + dta["id"]
        + " - "
        + dta["text_no_newline"]
    )
    output_string = dta.output_string.str.cat(sep="\n\n")
    dtst_path.write_text(output_string)


def create_output_string(dta):
    dta["output_string"] = (
        dta["id"]
        + " - "
        + dta["year_month"]
        + " - "
        + dta["comment"]
        + " - "
        + dta["text_combined"]
    )
    return dta.output_string.str.cat(sep="\n")


def load_model(modl_path):
    model_bytes = modl_path.read_bytes()
    return tp.LDAModel.loads(model_bytes)


def create_lda_output_files(
    target_label: str, topic_model_path: Path, dta: pd.DataFrame, num_topics: int, model_type="LDA"
):
    if not model_type in ["LDA", "PT"]:
        raise ValueError("model_type must be either LDA or PT")
    topic_model_num_path = topic_model_path / f"{target_label}_{str(num_topics).zfill(4)}_topics"
    topic_model_num_path.mkdir(parents=True, exist_ok=True)
    target_num_label = f"{target_label}_{str(num_topics).zfill(4)}"
    modl_path = topic_model_num_path / f"{target_num_label}_{model_type}_model.modl"
    mdl = load_model(modl_path)
    pyvs_path = topic_model_num_path / f"{target_num_label}_{model_type}_pyLDAvis.html"
    create_pyLDAVis(mdl, pyvs_path)
    topt_path = topic_model_num_path / f"{target_num_label}_{model_type}_topics_terms.txt"
    create_topic_summery_text_page(mdl, topt_path)

    # domt_path = topic_model_num_path / f"{target_num_label}_dominate_topics.parquet"
    # domt_path = topic_model_num_path / f"{target_num_label}_dominate_topics.parquet"
    # dtxt_path = topic_model_num_path / f"{target_num_label}_dominate_topics.txt"
    # dtst_path = topic_model_num_path / f"{target_num_label}_dominate_topics_sample.txt"
    # tpex_path = topic_model_num_path / f"{target_num_label}_dominant_topics_terms.xlsx"
    # tgex_path = topic_model_num_path / f"{target_num_label}_pub_count_by_year.xlsx"
    ltxp_path = topic_model_path / f"{target_label}_tomotopy_lemmatized_texts.parquet"
    dtdx_path = topic_model_num_path / f"{target_num_label}_dominant_topic_by_doc.xlsx"
    dtcx_path = topic_model_num_path / f"{target_num_label}_dominant_topic_count.xlsx"
    dtdp_path = topic_model_num_path / f"{target_num_label}_dominant_topic_by_doc.parquet"
    corp_path = topic_model_path / f"{target_label}_corpus.tomotopy"
    extract_dominant_topic(mdl, ltxp_path, dtdx_path, corp_path, dtdp_path, dtcx_path)
    # dta = create_dominant_topic_dta(domt_path, dta, text_dom_topics, percents_dom_topic)
    # create_dominant_topic_text_file(dta, dtxt_path, tpex_path, tgex_path)
    # create_top_sample_dominant_topic_text_file(dta, dtst_path)
