# -*- coding: utf-8 -*-

from os import makedirs
from os.path import exists, join
from time import time
import spacy
import pandas as pd
import gc

from constants import (
    VOC_PATH, TEXT, LEMMA, IWNLP, POS, TOK_IDX, SENT_START, ENT_IOB, ENT_TYPE, ENT_IDX,
    TOKEN, SENT_IDX, HASH, NOUN_PHRASE, NLP_PATH, PUNCT, SPACE, NUM, DET, TITLE, DESCR, ETL_PATH
)
from nlp_lemmatizer_plus import LemmatizerPlus
from utils import tprint

FIELDS = [HASH, TOK_IDX, SENT_IDX, TEXT, TOKEN, POS, ENT_IOB, ENT_IDX, ENT_TYPE, NOUN_PHRASE]


class NLPProcessor(object):

    def __init__(self, spacy_path, lemmatizer_path='../data/IWNLP.Lemmatizer_20170501.json', logg=None):
        self.logg = logg if logg else print

        # ------ load spacy and iwnlp ------
        logg("loading spacy")
        self.nlp = spacy.load(spacy_path)  # <-- load with dependency parser (slower)
        # nlp = spacy.load(de, disable=['parser'])

        if exists(VOC_PATH):
            logg("reading vocab from " + VOC_PATH)
            self.nlp.vocab.from_disk(VOC_PATH)

        logg("loading IWNLPWrapper")
        self.lemmatizer = LemmatizerPlus(lemmatizer_path, self.nlp)
        self.nlp.add_pipe(self.lemmatizer)
        self.stringstore = self.nlp.vocab.strings

    def read_process_store(self, file_path, corpus_name, store=True, vocab_to_disk=False,
                           start=0, stop=None, **kwargs):
        logg = self.logg
        logg("*** start new corpus: " + corpus_name)
        t0 = time()

        # read the etl dataframe
        slicing = "[{:d}:{:d}]".format(start, stop) if (start or stop) else ''
        logg("{}: reading corpus{} from {}".format(corpus_name, slicing, file_path))
        df = self.read(file_path, start=start, stop=stop)

        logg('collect: %d' % gc.collect())

        # start the nlp pipeline
        logg(corpus_name + ": start processing")
        # self.check_docs(df); return
        df = self.process_docs(df)
        logg('collect: %d' % gc.collect())

        if kwargs.get('print', False):
            # print dataframe
            tprint(df, kwargs.get('head', 10))
        if store:
            # store dataframe, free memory first
            if start or stop:
                suffix = '_{:d}_{:d}_nlp'.format(start, stop-1)
            else:
                suffix = '_nlp'
            self.store(corpus_name, df, suffix=suffix)
        if vocab_to_disk:
            # stored with each corpus, in case anythings goes wrong
            logg("writing spacy vocab to disk: " + VOC_PATH)
            # self.nlp.to_disk(SPACY_PATH)
            makedirs(VOC_PATH, exist_ok=True)
            self.nlp.vocab.to_disk(VOC_PATH)

        t1 = int(time() - t0)
        logg("{:s}: done in {:02d}:{:02d}:{:02d}".format(corpus_name, t1//3600, (t1//60) % 60, t1 % 60))

    def check_docs(self, text_df):
        for i, kv in enumerate(text_df.itertuples()):
            key, title, descr, text = kv
            doc = self.nlp('\n'.join(filter(None, [title, descr, text])))
            for token in doc:
                print(token.i, token.is_sent_start, token.text)

    def process_docs(self, text_df, steps=100):
        """ main function for sending the dataframes from the ETL pipeline to the NLP pipeline """
        steps = max(min(steps, len(text_df)), 1)
        step_len = max(100//steps, 1)
        percent = max(len(text_df) // steps, 1)
        chunk_idx = done = 0
        docs = []

        # process each doc in corpus
        for i, kv in enumerate(text_df.itertuples()):
            # log progress
            if i % percent == 0:
                if i > 0:
                    self.logg("  {:d}%: {:d} documents processed".format(done, i))
                done += step_len

            key, title, descr, text = kv
            # build spacy doc
            doc = self.nlp('\n'.join(filter(None, [title, descr, text])))

            # annotated phrases
            noun_phrases = dict()
            for chunk in doc.noun_chunks:
                chunk_idx += 1
                for token in chunk:
                    noun_phrases[token.i] = chunk_idx

            # extract relevant attributes
            attr = [[key,
                     str(token.text), str(token.lemma_), token._.iwnlp_lemmas, str(token.pos_),
                     int(token.i), 1 if (token.is_sent_start or token.i == 0) else 0,
                     token.ent_iob_, token.ent_type_,
                     noun_phrases.get(token.i, 0)
                     ] for token in doc]
            # add list of token to all docs
            docs += attr

        return self.df_from_docs(docs)

    def store(self, corpus, df, suffix=''):
        """returns the file path where the dataframe was stores"""
        makedirs(NLP_PATH, exist_ok=True)
        fname = join(NLP_PATH, corpus + suffix + '.pickle')
        self.logg(corpus + ': saving to ' + fname)
        df.to_pickle(fname)

    def read(self, f, start=0, stop=None):
        """ reads a dataframe from pickle format """
        df = pd.read_pickle(f)[[TITLE, DESCR, TEXT]].iloc[start:stop]
        # lazy hack for dewiki_new
        if 'dewiki' in f:
            goodids = pd.read_pickle(join(ETL_PATH, 'dewiki_good_ids.pickle'))
            df = df[df.index.isin(goodids.index)]
        self.logg('using {:d} documents'.format(len(df)))
        return df.copy()

    @staticmethod
    def df_from_docs(docs):
        """
        Creates a pandas DataFrame from a given spacy.doc that contains only nouns and noun phrases.
        :param docs: list of tokens (tuples with attributes) from spacy.doc
        :return:    pandas.DataFrame
        """
        df = pd.DataFrame(docs,
                          columns=[HASH, TEXT, LEMMA, IWNLP, POS,
                                   TOK_IDX, SENT_START, ENT_IOB, ENT_TYPE, NOUN_PHRASE])
        # create Tokens from IWNLP lemmatization, else from spacy lemmatization (or original text)
        mask_iwnlp = ~df[IWNLP].isnull()
        df.loc[mask_iwnlp, TOKEN] = df.loc[mask_iwnlp, IWNLP]
        df.loc[~mask_iwnlp, TOKEN] = df.loc[~mask_iwnlp, LEMMA]
        # fixes wrong POS tagging for punctuation (TODO: untested at this step)
        mask_punct = df[TOKEN].isin(list('[]<>/–%'))
        df.loc[mask_punct, POS] = PUNCT
        # set an index for each sentence
        df[SENT_IDX] = df[SENT_START].cumsum()
        # set an index for each entity
        df[ENT_IDX] = (df[ENT_IOB] == 'B')
        df[ENT_IDX] = df[ENT_IDX].cumsum()
        df.loc[df[ENT_IOB] == 'O', ENT_IDX] = 0
        # convert for space efficiency
        df[POS] = df[POS].astype("category")
        df[ENT_IOB] = df[ENT_IOB].astype("category")
        df[ENT_TYPE] = df[ENT_TYPE].astype("category")
        return df[FIELDS]

# -*- coding: utf-8 -*-

from os import makedirs
from os.path import exists, join
from time import time
import spacy
import pandas as pd
import gc

from constants import (
    VOC_PATH, TEXT, LEMMA, IWNLP, POS, TOK_IDX, SENT_START, ENT_IOB, ENT_TYPE, ENT_IDX,
    TOKEN, SENT_IDX, HASH, NOUN_PHRASE, NLP_PATH, PUNCT, SPACE, NUM, DET, TITLE, DESCR, ETL_PATH
)
from nlp_lemmatizer_plus import LemmatizerPlus
from utils import tprint

FIELDS = [HASH, TOK_IDX, SENT_IDX, TEXT, TOKEN, POS, ENT_IOB, ENT_IDX, ENT_TYPE, NOUN_PHRASE]


class NLPProcessor(object):

    def __init__(self, spacy_path, lemmatizer_path='../data/IWNLP.Lemmatizer_20170501.json', logg=None):
        self.logg = logg if logg else print

        # ------ load spacy and iwnlp ------
        logg("loading spacy")
        self.nlp = spacy.load(spacy_path)  # <-- load with dependency parser (slower)
        # nlp = spacy.load(de, disable=['parser'])

        if exists(VOC_PATH):
            logg("reading vocab from " + VOC_PATH)
            self.nlp.vocab.from_disk(VOC_PATH)

        logg("loading IWNLPWrapper")
        self.lemmatizer = LemmatizerPlus(lemmatizer_path, self.nlp)
        self.nlp.add_pipe(self.lemmatizer)
        self.stringstore = self.nlp.vocab.strings

    def read_process_store(self, file_path, corpus_name, store=True, vocab_to_disk=False,
                           start=0, stop=None, **kwargs):
        logg = self.logg
        logg("*** start new corpus: " + corpus_name)
        t0 = time()

        # read the etl dataframe
        slicing = "[{:d}:{:d}]".format(start, stop) if (start or stop) else ''
        logg("{}: reading corpus{} from {}".format(corpus_name, slicing, file_path))
        df = self.read(file_path, start=start, stop=stop)

        logg('collect: %d' % gc.collect())

        # start the nlp pipeline
        logg(corpus_name + ": start processing")
        # self.check_docs(df); return
        df = self.process_docs(df)
        logg('collect: %d' % gc.collect())

        if kwargs.get('print', False):
            # print dataframe
            tprint(df, kwargs.get('head', 10))
        if store:
            # store dataframe, free memory first
            if start or stop:
                suffix = '_{:d}_{:d}_nlp'.format(start, stop-1)
            else:
                suffix = '_nlp'
            self.store(corpus_name, df, suffix=suffix)
        if vocab_to_disk:
            # stored with each corpus, in case anythings goes wrong
            logg("writing spacy vocab to disk: " + VOC_PATH)
            # self.nlp.to_disk(SPACY_PATH)
            makedirs(VOC_PATH, exist_ok=True)
            self.nlp.vocab.to_disk(VOC_PATH)

        t1 = int(time() - t0)
        logg("{:s}: done in {:02d}:{:02d}:{:02d}".format(corpus_name, t1//3600, (t1//60) % 60, t1 % 60))

    def check_docs(self, text_df):
        for i, kv in enumerate(text_df.itertuples()):
            key, title, descr, text = kv
            doc = self.nlp('\n'.join(filter(None, [title, descr, text])))
            for token in doc:
                print(token.i, token.is_sent_start, token.text)

    def process_docs(self, text_df, steps=100):
        """ main function for sending the dataframes from the ETL pipeline to the NLP pipeline """
        steps = max(min(steps, len(text_df)), 1)
        step_len = max(100//steps, 1)
        percent = max(len(text_df) // steps, 1)
        chunk_idx = done = 0
        docs = []

        # process each doc in corpus
        for i, kv in enumerate(text_df.itertuples()):
            # log progress
            if i % percent == 0:
                if i > 0:
                    self.logg("  {:d}%: {:d} documents processed".format(done, i))
                done += step_len

            key, title, descr, text = kv
            # build spacy doc
            doc = self.nlp('\n'.join(filter(None, [title, descr, text])))

            # annotated phrases
            noun_phrases = dict()
            for chunk in doc.noun_chunks:
                chunk_idx += 1
                for token in chunk:
                    noun_phrases[token.i] = chunk_idx

            # extract relevant attributes
            attr = [[key,
                     str(token.text), str(token.lemma_), token._.iwnlp_lemmas, str(token.pos_),
                     int(token.i), 1 if (token.is_sent_start or token.i == 0) else 0,
                     token.ent_iob_, token.ent_type_,
                     noun_phrases.get(token.i, 0)
                     ] for token in doc]
            # add list of token to all docs
            docs += attr

        return self.df_from_docs(docs)

    def store(self, corpus, df, suffix=''):
        """returns the file path where the dataframe was stores"""
        makedirs(NLP_PATH, exist_ok=True)
        fname = join(NLP_PATH, corpus + suffix + '.pickle')
        self.logg(corpus + ': saving to ' + fname)
        df.to_pickle(fname)

    def read(self, f, start=0, stop=None):
        """ reads a dataframe from pickle format """
        df = pd.read_pickle(f)[[TITLE, DESCR, TEXT]].iloc[start:stop]
        # lazy hack for dewiki_new
        if 'dewiki' in f:
            goodids = pd.read_pickle(join(ETL_PATH, 'dewiki_good_ids.pickle'))
            df = df[df.index.isin(goodids.index)]
        self.logg('using {:d} documents'.format(len(df)))
        return df.copy()

    @staticmethod
    def df_from_docs(docs):
        """
        Creates a pandas DataFrame from a given spacy.doc that contains only nouns and noun phrases.
        :param docs: list of tokens (tuples with attributes) from spacy.doc
        :return:    pandas.DataFrame
        """
        df = pd.DataFrame(docs,
                          columns=[HASH, TEXT, LEMMA, IWNLP, POS,
                                   TOK_IDX, SENT_START, ENT_IOB, ENT_TYPE, NOUN_PHRASE])
        # create Tokens from IWNLP lemmatization, else from spacy lemmatization (or original text)
        mask_iwnlp = ~df[IWNLP].isnull()
        df.loc[mask_iwnlp, TOKEN] = df.loc[mask_iwnlp, IWNLP]
        df.loc[~mask_iwnlp, TOKEN] = df.loc[~mask_iwnlp, LEMMA]
        # fixes wrong POS tagging for punctuation (TODO: untested at this step)
        mask_punct = df[TOKEN].isin(list('[]<>/–%'))
        df.loc[mask_punct, POS] = PUNCT
        # set an index for each sentence
        df[SENT_IDX] = df[SENT_START].cumsum()
        # set an index for each entity
        df[ENT_IDX] = (df[ENT_IOB] == 'B')
        df[ENT_IDX] = df[ENT_IDX].cumsum()
        df.loc[df[ENT_IOB] == 'O', ENT_IDX] = 0
        # convert for space efficiency
        df[POS] = df[POS].astype("category")
        df[ENT_IOB] = df[ENT_IOB].astype("category")
        df[ENT_TYPE] = df[ENT_TYPE].astype("category")
        return df[FIELDS]

# -*- coding: utf-8 -*- # ------ load spacy and iwnlp ------ # <-- load with dependency parser (slower) # nlp = spacy.load(de, disable=['parser']) # read the etl dataframe # start the nlp pipeline # self.check_docs(df); return # print dataframe # store dataframe, free memory first # stored with each corpus, in case anythings goes wrong # self.nlp.to_disk(SPACY_PATH) main function for sending the dataframes from the ETL pipeline to the NLP pipeline # process each doc in corpus # log progress # build spacy doc # annotated phrases # extract relevant attributes # add list of token to all docs returns the file path where the dataframe was stores reads a dataframe from pickle format # lazy hack for dewiki_new Creates a pandas DataFrame from a given spacy.doc that contains only nouns and noun phrases.         :param docs: list of tokens (tuples with attributes) from spacy.doc         :return:    pandas.DataFrame # create Tokens from IWNLP lemmatization, else from spacy lemmatization (or original text) # fixes wrong POS tagging for punctuation (TODO: untested at this step) # set an index for each sentence # set an index for each entity # convert for space efficiency