import pandas as pd
from tqdm import tqdm
import pickle
from multiprocessing import Pool
import multiprocessing
import textacy
import textacy.preprocessing
import textacy.representations
import textacy.tm
from functools import partial
import warnings

# import en_core_web_sm
# import fr_core_news_lg

tqdm.pandas()
# lang = en_core_web_sm.load()

warnings.simplefilter(action="ignore", category=FutureWarning)

preproc = textacy.preprocessing.make_pipeline(
    textacy.preprocessing.normalize.unicode,
    textacy.preprocessing.normalize.bullet_points,
    textacy.preprocessing.normalize.quotation_marks,
    textacy.preprocessing.normalize.whitespace,
    textacy.preprocessing.normalize.hyphenated_words,
    textacy.preprocessing.remove.brackets,
    textacy.preprocessing.replace.currency_symbols,
    textacy.preprocessing.remove.html_tags,
)


def from_dict_to_frame(indexed_dict):
    data = {k: [v] for k, v in indexed_dict.items()}
    df = pd.DataFrame.from_dict(data).T
    df.columns = ["text"]
    df = df.explode("text")
    return df


def extract_terms_df(
    data,
    text_var,
    index_var,
    ngs=True,
    ents=True,
    ncs=False,
    multiprocess=True,
    sample_size=100000,
    drop_emoji=True,
    ngrams=(2, 2),
    remove_punctuation=True,
    include_pos=["NOUN", "PROPN", "ADJ"],
    include_types=["PERSON", "ORG"],
    language="en",
):

    if language == "zh":
        # nlp = zh_core_web_sm.load()
        lang = textacy.load_spacy_lang("zh_core_web_sm", disable=())
    elif language == "en":
        # nlp = en_core_web_sm.load()
        lang = textacy.load_spacy_lang("en_core_web_sm", disable=())
    elif language == "fr":
        # nlp = fr_core_news_lg.load()
        lang = textacy.load_spacy_lang("fr_core_news_lg", disable=())
    else:
        pass

    def extract_terms(
        tuple,  # (index, text)
        ngs=True,
        ents=True,
        ncs=False,
        ngrams=(2, 2),
        drop_emoji=True,
        remove_punctuation=False,
        include_pos=["NOUN", "PROPN", "ADJ"],
        include_types=["PERSON", "ORG"],
    ):

        index = tuple[0]
        text = tuple[1]

        prepro_text = preproc(str(text))
        if drop_emoji == True:
            prepro_text = textacy.preprocessing.replace.emojis(prepro_text, repl="")

        if remove_punctuation == True:
            prepro_text = textacy.preprocessing.remove.punctuation(prepro_text)

        doc = textacy.make_spacy_doc(prepro_text, lang=lang)

        terms = []

        if ngs:
            ngrams_terms = list(
                textacy.extract.terms(
                    doc,
                    ngs=partial(
                        textacy.extract.ngrams,
                        n=ngrams,
                        filter_punct=True,
                        filter_stops=True,
                        include_pos=include_pos,
                    ),
                    dedupe=False,
                )
            )

            terms.append(ngrams_terms)

        if ents:
            ents_terms = list(
                textacy.extract.terms(
                    doc,
                    ents=partial(textacy.extract.entities, include_types=include_types),
                    dedupe=False,
                )
            )
            terms.append(ents_terms)

        if ncs:
            ncs_terms = list(
                textacy.extract.terms(
                    doc,
                    ncs=partial(textacy.extract.noun_chunks, drop_determiners=True),
                    dedupe=False,
                )
            )

            noun_chunks = [x for x in ncs_terms if len(x) >= 3]
            terms.append(noun_chunks)

        final = [item for sublist in terms for item in sublist]
        final = list(set(final))

        df = [
            (term.text, term.lemma_.lower(), term.label_, term.__len__())
            for term in final
        ]
        df = pd.DataFrame(df, columns=["text", "lemma", "ent", "ngrams"])
        df["text_index"] = index

        return df

    """
    This function extracts terms from a column in a DataFrame. It can extract in a multiprocessing way
    It outputs a dataframe with the list of terms and a table with the indexed terms
    """

    data = data[data[text_var].notna()]
    data = data.sample(min(sample_size, len(data)))

    sentences = data[text_var].to_list()
    indexes = data[index_var].to_list()
    inputs = [(x, y) for x, y in zip(indexes, sentences)]

    """if multiprocess:
        with Pool(multiprocessing.cpu_count() - 2) as p:
            res = list(
                tqdm(
                    p.imap(
                        partial(
                            extract_terms,
                            ngs=ngs,
                            ents=ents,
                            ncs=ncs,
                            drop_emoji=drop_emoji,
                            remove_punctuation=remove_punctuation,
                            ngrams=ngrams,
                            include_pos=include_pos,
                            include_types=include_types,
                        ),
                        inputs,
                    ),
                    total=len(inputs),
                )
            )

        final_res = pd.concat([x for x in res])"""

    # else:
    res = list(
        tqdm(
            map(
                partial(
                    extract_terms,
                    ngs=ngs,
                    ents=ents,
                    ncs=ncs,
                    drop_emoji=drop_emoji,
                    remove_punctuation=remove_punctuation,
                    ngrams=ngrams,
                    include_pos=include_pos,
                    include_types=include_types,
                ),
                inputs,
            ),
            total=len(inputs),
        )
    )

    final_res = pd.concat([x for x in res])

    terms = (
        final_res.groupby(["text", "lemma", "ent", "ngrams"])
        .agg(count_terms=("text_index", "count"))
        .reset_index()
    )

    # duplicates to get rid of
    terms = terms.sort_values(["text", "ent"]).reset_index()
    terms = terms.drop_duplicates(["text"], keep="first")
    terms = terms.sort_values("count_terms", ascending=False)
    terms = terms.set_index("text")

    terms_indexed = final_res[["text", "text_index"]].drop_duplicates()
    terms_indexed = terms_indexed.rename(columns={"text_index": index_var})

    terms_indexed = terms_indexed.groupby(index_var)["text"].apply(list)

    return terms, terms_indexed
