# coding: utf8
"""Conll parser"""
from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from __future__ import print_function

import re
import sys
import codecs
import argparse
import time
import os
import io
import pickle

import spacy
from spacy.tokens import Doc

import numpy as np

from tqdm import tqdm

from .compat import unicode_
from .document import Mention, Document, Speaker, EmbeddingExtractor, MISSING_WORD
from .utils import parallel_process

PACKAGE_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
REMOVED_CHAR = ["/", "%", "*"]
NORMALIZE_DICT = {"/.": ".",
                  "/?": "?",
                  "-LRB-": "(",
                  "-RRB-": ")",
                  "-LCB-": "{",
                  "-RCB-": "}",
                  "-LSB-": "[",
                  "-RSB-": "]"}

CONLL_GENRES = {"bc": 0, "bn": 1, "mz": 2, "nw": 3, "pt": 4, "tc": 5, "wb": 6}

FEATURES_NAMES = ["mentions_features",          # 0
                  "mentions_labels",            # 1
                  "mentions_pairs_length",      # 2
                  "mentions_pairs_start_index", # 3
                  "mentions_spans",             # 4
                  "mentions_words",             # 5
                  "pairs_ant_index",            # 6
                  "pairs_features",             # 7
                  "pairs_labels",               # 8
                  "locations",                  # 9
                  "conll_tokens",               # 10
                  "spacy_lookup",               # 11
                  "doc",                        # 12
                  ]

MISSED_MENTIONS_FILE = os.path.join(PACKAGE_DIRECTORY, "test_mentions_identification.txt")
SENTENCES_PATH = os.path.join(PACKAGE_DIRECTORY, "test_sentences.txt")

###################
### UTILITIES #####

def clean_token(token):
    cleaned_token = token
    if cleaned_token in NORMALIZE_DICT:
        cleaned_token = NORMALIZE_DICT[cleaned_token]
    if cleaned_token not in REMOVED_CHAR:
        for char in REMOVED_CHAR:
            cleaned_token = cleaned_token.replace(char, u'')
    if len(cleaned_token) == 0:
        cleaned_token = ","
    return cleaned_token

def mention_words_idx(embed_extractor, mention, debug=False):
    # index of the word in the tuned embeddings no need for normalizing,
    # it is already performed in set_mentions_features()
    # We take them in the tuned vocabulary which is a smaller voc tailored from conll
    words = []
    for _, w in sorted(mention.words_embeddings_.items()):
        if w not in embed_extractor.tun_idx:
            if debug: print("No matching tokens in tuned voc for word ", w, "surrounding or inside mention", mention)
            words.append(MISSING_WORD)
        else:
            words.append(w)
    return [embed_extractor.tun_idx[w] for w in words]

def check_numpy_array(feature, array, n_mentions_list, compressed=True):
    for n_mentions in n_mentions_list:
        if feature == FEATURES_NAMES[0]:
            assert array.shape[0] == len(n_mentions)
            if compressed:
                assert np.array_equiv(array[:, 3], np.array([len(n_mentions)] * len(n_mentions)))
                assert np.max(array[:, 2]) == len(n_mentions)-1
                assert np.min(array[:, 2]) == 0
        elif feature == FEATURES_NAMES[1]:
            assert array.shape[0] == len(n_mentions)
        elif feature == FEATURES_NAMES[2]:
            assert array.shape[0] == len(n_mentions)
            assert np.array_equiv(array[:, 0], np.array(list(range(len(n_mentions)))))
        elif feature == FEATURES_NAMES[3]:
            assert array.shape[0] == len(n_mentions)
            assert np.array_equiv(array[:, 0], np.array([p*(p-1)/2 for p in range(len(n_mentions))]))
        elif feature == FEATURES_NAMES[4]:
            assert array.shape[0] == len(n_mentions)
        elif feature == FEATURES_NAMES[5]:
            assert array.shape[0] == len(n_mentions)
        elif feature == FEATURES_NAMES[6]:
            assert array.shape[0] == len(n_mentions)*(len(n_mentions)-1)/2
            assert np.max(array) == len(n_mentions)-2
        elif feature == FEATURES_NAMES[7]:
            if compressed:
                assert array.shape[0] == len(n_mentions)*(len(n_mentions)-1)/2
                assert np.max(array[:, 7]) == len(n_mentions)-2
                assert np.min(array[:, 7]) == 0
        elif feature == FEATURES_NAMES[8]:
            assert array.shape[0] == len(n_mentions)*(len(n_mentions)-1)/2

###############################################################################################
### PARALLEL FCT (has to be at top-level of the module to be pickled for multiprocessing) #####
def load_file(full_name, debug=False):
    '''
    load a *._conll file
    Input: full_name: path to the file
    Output: list of tuples for each conll doc in the file, where the tuple contains:
        (utts_text ([str]): list of the utterances in the document 
         utts_tokens ([[str]]): list of the tokens (conll words) in the document 
         utts_corefs: list of coref objects (dicts) with the following properties:
            coref['label']: id of the coreference cluster,
            coref['start']: start index (index of first token in the utterance),
            coref['end': end index (index of last token in the utterance).
         utts_speakers ([str]): list of the speaker associated to each utterances in the document 
         name (str): name of the document
         part (str): part of the document
        )
    '''
    docs = []
    with io.open(full_name, 'rt', encoding='utf-8', errors='strict') as f:
        lines = list(f)#.readlines()
        utts_text = []
        utts_tokens = []
        utts_corefs = []
        utts_speakers = []
        tokens = []
        corefs = []
        index = 0
        speaker = ""
        name = ""
        part = ""
        for li, line in enumerate(lines):
            cols = line.split()
            if debug: print("line", li, "cols:", cols)
            # End of utterance
            if len(cols) == 0:
                if tokens:
                    if debug: print("End of utterance")
                    utts_text.append(u''.join(t + u' ' for t in tokens))
                    utts_tokens.append(tokens)
                    utts_speakers.append(speaker)
                    utts_corefs.append(corefs)
                    tokens = []
                    corefs = []
                    index = 0
                    speaker = ""
                    continue
            # End of doc
            elif len(cols) == 2:
                if debug: print("End of doc")
                if cols[0] == "#end":
                    if debug: print("Saving doc")
                    docs.append((utts_text, utts_tokens, utts_corefs, utts_speakers, name, part))
                    utts_text = []
                    utts_tokens = []
                    utts_corefs = []
                    utts_speakers = []
                else:
                    raise ValueError("Error on end line " + line)
            # New doc
            elif len(cols) == 5:
                if debug: print("New doc")
                if cols[0] == "#begin":
                    name = re.match(r"\((.*)\);", cols[2]).group(1)
                    try:
                        part = cols[4]
                    except ValueError:
                        print("Error parsing document part " + line)
                    if debug: print("New doc", name, part, name[:2])
                    tokens = []
                    corefs = []
                    index = 0
                else:
                    raise ValueError("Error on begin line " + line)
            # Inside utterance
            elif len(cols) > 7:
                if debug: print("Inside utterance")
                assert (cols[0] == name and int(cols[1]) == int(part)), "Doc name or part error " + line
                assert (int(cols[2]) == index), "Index error on " + line
                if speaker:
                    assert (cols[9] == speaker), "Speaker changed in " + line + speaker
                else:
                    speaker = cols[9]
                    if debug: print("speaker", speaker)
                if cols[-1] != u'-':
                    coref_expr = cols[-1].split(u'|')
                    if debug: print("coref_expr", coref_expr)
                    if not coref_expr:
                        raise ValueError("Coref expression empty " + line)
                    for tok in coref_expr:
                        if debug: print("coref tok", tok)
                        try:
                            match = re.match(r"^(\(?)(\d+)(\)?)$", tok)
                        except:
                            print("error getting coreferences for line " + line)
                        assert match is not None, "Error parsing coref " + tok + " in " + line
                        num = match.group(2)
                        assert (num is not u''), "Error parsing coref " + tok + " in " + line
                        if match.group(1) == u'(':
                            if debug: print("New coref", num)
                            corefs.append({'label': num, 'start': index, 'end': None})
                        if match.group(3) == u')':
                            j = None
                            for i in range(len(corefs)-1, -1, -1):
                                if debug: print("i", i)
                                if corefs[i]['label'] == num and corefs[i]['end'] is None:
                                    j = i
                                    break
                            assert (j is not None), "coref closing error " + line
                            if debug: print("End coref", num)
                            corefs[j]['end'] = index
                tokens.append(clean_token(cols[3]))
                index += 1
            else:
                raise ValueError("Line not standard " + line)
    return docs

def set_feats(doc):
    doc.set_mentions_features()

def get_feats(doc, i):
    return doc.get_feature_array(doc_id=i)

def gather_feats(gathering_array, array, feat_name, pairs_ant_index, pairs_start_index):
    if gathering_array is None:
        gathering_array = array
    else:
        if feat_name == FEATURES_NAMES[6]:
            array = [a + pairs_ant_index for a in array]
        elif feat_name == FEATURES_NAMES[3]:
            array = [a + pairs_start_index for a in array]
        gathering_array += array
    return feat_name, gathering_array

def read_file(full_name):
    doc = ""
    with io.open(full_name, 'rt', encoding='utf-8', errors='strict') as f:
        doc = f.read()
    return doc

###################
### ConllDoc #####

class ConllDoc(Document):
    def __init__(self, name, part, *args, **kwargs):
        self.name = name
        self.part = part
        self.feature_matrix = {}
        self.conll_tokens = []
        self.conll_lookup = []
        self.gold_corefs = []
        self.missed_gold = []
        super(ConllDoc, self).__init__(*args, **kwargs)

    def get_conll_spacy_lookup(self, conll_tokens, spacy_tokens, debug=False):
        '''
        Compute a look up table between spacy tokens (from spacy tokenizer)
        and conll pre-tokenized tokens
        Output: list[conll_index] => list of associated spacy tokens (assume spacy tokenizer has a finer granularity)
        '''
        lookup = []
        c_iter = (t for t in conll_tokens)
        s_iter = enumerate(t for t in spacy_tokens)
        i, s_tok = next(s_iter)
        for c_tok in c_iter:
            #if debug: print("conll", c_tok, "spacy", s_tok, "index", i)
            c_lookup = []
            while i is not None and len(c_tok) and c_tok.startswith(s_tok.text):
                c_lookup.append(i)
                c_tok = c_tok[len(s_tok):]
                i, s_tok = next(s_iter, (None, None))
                if debug and len(c_tok): print("eating token: conll", c_tok, "spacy", s_tok, "index", i)
            assert len(c_lookup), "Unmatched conll and spacy tokens"
            lookup.append(c_lookup)
        return lookup

    def add_conll_utterance(self, parsed, tokens, corefs, speaker_id, use_gold_mentions=False, debug=False):
        conll_lookup = self.get_conll_spacy_lookup(tokens, parsed)
        self.conll_tokens.append(tokens)
        self.conll_lookup.append(conll_lookup)
        # Convert conll tokens coref index in spacy tokens indexes
        identified_gold = [False] * len(corefs)
        for coref in corefs:
            assert (coref['label'] is not None and coref['start'] is not None and coref['end'] is not None), \
                ("Error in coreference " + coref + " in " + parsed)
            coref['start'] = conll_lookup[coref['start']][0]
            coref['end'] = conll_lookup[coref['end']][-1]

        if speaker_id not in self.speakers:
            speaker_name = speaker_id.split(u'_')
            if debug: print("New speaker: ", speaker_id, "name: ", speaker_name)
            self.speakers[speaker_id] = Speaker(speaker_id, speaker_name)
        if use_gold_mentions:
            for coref in corefs:
                #            print("coref['label']", coref['label'])
                #            print("coref text",parsed[coref['start']:coref['end']])
                mention = Mention(parsed[coref['start']:coref['end']], len(self.mentions), len(self.utterances),
                                  self.n_sents, speaker=self.speakers[speaker_id], gold_label=coref['label'])
                self.mentions.append(mention)
                #            print("mention: ", mention, "label", mention.gold_label)
        else:
            self._extract_mentions(parsed, len(self.utterances), self.n_sents, self.speakers[speaker_id])
            # Assign a gold label to mentions which have one
            if debug: print("Check corefs", corefs)
            for i, coref in enumerate(corefs):
                for m in self.mentions:
                    if m.utterance_index != len(self.utterances):
                        continue
                    # if debug: print("Checking mention", m, m.utterance_index, m.start, m.end)
                    if coref['start'] == m.start and coref['end'] == m.end - 1:
                        m.gold_label = coref['label']
                        identified_gold[i] = True
                        # if debug: print("Gold mention found:", m, coref['label'])
            for found, coref in zip(identified_gold, corefs):
                if not found:
                    self.missed_gold.append([self.name, self.part, str(len(self.utterances)), parsed.text, parsed[coref['start']:coref['end']+1].text])
                    if debug:
                        print("❄️ gold mention not in predicted mentions", coref, parsed[coref['start']:coref['end']+1])
        self.utterances.append(parsed)
        self.gold_corefs.append(corefs)
        self.utterances_speaker.append(self.speakers[speaker_id])
        self.n_sents += len(list(parsed.sents))

    def get_single_mention_features_conll(self, mention, compressed=True):
        ''' Compressed or not single mention features'''
        if not compressed:
            _, features = self.get_single_mention_features(mention)
            return features[np.newaxis, :]
        feat_l = [mention.features_["01_MentionType"],
                  mention.features_["02_MentionLength"],
                  mention.index,
                  len(self.mentions),
                  mention.features_["04_IsMentionNested"],
                  self.genre_,
                 ]
        return feat_l

    def get_pair_mentions_features_conll(self, m1, m2, compressed=True):
        ''' Compressed or not single mention features'''
        if not compressed:
            _, features = self.get_pair_mentions_features(m1, m2)
            return features[np.newaxis, :]
        features_, _ = self.get_pair_mentions_features(m1, m2)
        feat_l = [features_["00_SameSpeaker"],
                  features_["01_AntMatchMentionSpeaker"],
                  features_["02_MentionMatchSpeaker"],
                  features_["03_HeadsAgree"],
                  features_["04_ExactStringMatch"],
                  features_["05_RelaxedStringMatch"],
                  features_["06_SentenceDistance"],
                  features_["07_MentionDistance"],
                  features_["08_Overlapping"],
                 ]
        return feat_l

    def get_feature_array(self, doc_id, feature=None, compressed=True, debug=True):
        """
        Prepare feature array:
            mentions_spans: (N, S)
            mentions_words: (N, W)
            mentions_features: (N, Fs)
            mentions_labels: (N, 1)
            mentions_pairs_start_index: (N, 1) index of beggining of pair list in pair_labels
            mentions_pairs_length: (N, 1) number of pairs (i.e. nb of antecedents) for each mention 
            pairs_features: (P, Fp)
            pairs_labels: (P, 1)
            pairs_ant_idx: (P, 1) => indexes of antecedents mention for each pair (mention index in doc)
        """ 
        if not self.mentions:
            print("No mention in this doc !")
            return {}
        if debug: print("🛎 features matrices")
        mentions_spans = []
        mentions_words = []
        mentions_features = []
        pairs_ant_idx = []
        pairs_features = []
        pairs_labels = []
        mentions_labels = []
        mentions_pairs_start = []
        mentions_pairs_length = []
        mentions_location = []
        n_mentions = 0
        total_pairs = 0
        if debug: print("mentions", self.mentions, str([m.gold_label for m in self.mentions]))
        for mention_idx, antecedents_idx in list(self.get_candidate_pairs(max_distance=None, max_distance_with_match=None)):
            n_mentions += 1
            mention = self.mentions[mention_idx]
            mentions_spans.append(mention.spans_embeddings)
            w_idx = mention_words_idx(self.embed_extractor, mention)
            if w_idx is None:
                print("error in", self.name, self.part, mention.utterance_index)
            mentions_words.append(w_idx)
            mentions_features.append(self.get_single_mention_features_conll(mention, compressed))
            mentions_location.append([mention.start, mention.end, mention.utterance_index, mention_idx, doc_id])
            ants = [self.mentions[ant_idx] for ant_idx in antecedents_idx]
            no_antecedent = not any(ant.gold_label == mention.gold_label for ant in ants) or mention.gold_label is None
            if antecedents_idx:
                pairs_ant_idx += [idx for idx in antecedents_idx]
                pairs_features += [self.get_pair_mentions_features_conll(ant, mention, compressed) for ant in ants]
                ant_labels = [0 for ant in ants] if no_antecedent else [1 if ant.gold_label == mention.gold_label else 0 for ant in ants]
                pairs_labels += ant_labels
            mentions_labels.append(1 if no_antecedent else 0)
            mentions_pairs_start.append(total_pairs)
            total_pairs += len(ants)
            mentions_pairs_length.append(len(ants))

        out_dict = {FEATURES_NAMES[0]: mentions_features,
                    FEATURES_NAMES[1]: mentions_labels,
                    FEATURES_NAMES[2]: mentions_pairs_length,
                    FEATURES_NAMES[3]: mentions_pairs_start,
                    FEATURES_NAMES[4]: mentions_spans,
                    FEATURES_NAMES[5]: mentions_words,
                    FEATURES_NAMES[6]: pairs_ant_idx if pairs_ant_idx else None,
                    FEATURES_NAMES[7]: pairs_features if pairs_features else None,
                    FEATURES_NAMES[8]: pairs_labels if pairs_labels else None,
                    FEATURES_NAMES[9]: [mentions_location],
                    FEATURES_NAMES[10]: [self.conll_tokens],
                    FEATURES_NAMES[11]: [self.conll_lookup],
                    FEATURES_NAMES[12]: [{'name': self.name,
                                          'part': self.part,
                                          'utterances': list(str(u) for u in self.utterances),
                                          'mentions': list(str(m) for m in self.mentions)}],
                    }
        if debug:
            print("🚘 Summary")
            for k, v in out_dict.items():
                print(k, len(v))
        return n_mentions, total_pairs, out_dict

###################
### ConllCorpus #####
class ConllCorpus(object):
    def __init__(self, n_jobs=4, embed_path=PACKAGE_DIRECTORY+"/weights/", use_gold_mentions=False):
        self.n_jobs = n_jobs
        self.features = {}
        self.utts_text = []
        self.utts_tokens = []
        self.utts_corefs = []
        self.utts_speakers = []
        self.utts_doc_idx = []
        self.docs_names = []
        self.docs = []
        if embed_path is not None:
            self.embed_extractor = EmbeddingExtractor(embed_path)
        self.trainable_embed = []
        self.trainable_voc = []
        self.use_gold_mentions = use_gold_mentions

    def check_words_in_embeddings_voc(self, embedding, tuned=True, debug=False):
        print("🌋 Checking if words are in embedding voc")
        if tuned:
            embed_voc = embedding.tun_idx
        else:
            embed_voc = embedding.stat_idx
        missing_words = []
        missing_words_sents = []
        missing_words_doc = []
        for doc in self.docs:
            # if debug: print("Checking doc", doc.name, doc.part)
            for sent in doc.utterances:
                # if debug: print(sent.text)
                for word in sent:
                    w = embedding.normalize_word(word)
                    # if debug: print(w)
                    if w not in embed_voc:
                        missing_words.append(w)
                        missing_words_sents.append(sent.text)
                        missing_words_doc.append(doc.name + doc.part)
                        if debug:
                            out_str = "No matching tokens in tuned voc for " + w + \
                                      " in sentence " + sent.text + \
                                      " in doc " + doc.name + doc.part
                            print(out_str)
        return missing_words, missing_words_sents, missing_words_doc

    def test_sentences_words(self, save_file, debug=False):
        print("🌋 Saving sentence list")
        with io.open(save_file, "w", encoding='utf-8') as f:
            if debug: print("Sentences saved in", save_file)
            for doc in self.docs:
                out_str = "#begin document (" + doc.name + \
                          "); part " + doc.part + "\n"
                f.write(out_str)
                for sent in doc.utterances:
                    f.write(sent.text + '\n')
                out_str = "#end document\n\n"
                f.write(out_str)

    def save_sentences(self, save_file, debug=False):
        print("🌋 Saving sentence list")
        with io.open(save_file, "w", encoding='utf-8') as f:
            if debug: print("Sentences saved in", save_file)
            for doc in self.docs:
                out_str = "#begin document (" + doc.name + \
                          "); part " + doc.part + "\n"
                f.write(out_str)
                for sent in doc.utterances:
                    f.write(sent.text + '\n')
                out_str = "#end document\n\n"
                f.write(out_str)

    def build_key_file(self, data_path, key_file, debug=False):
        print("🌋 Building key file from corpus")
        print("Saving in", key_file)
        # Create a pool of processes. By default, one is created for each CPU in your machine.
        with io.open(key_file, "w", encoding='utf-8') as kf:
            if debug: print("Key file saved in", key_file)
            for dirpath, _, filenames in os.walk(data_path):
                print("In", dirpath)
                file_list = [os.path.join(dirpath, f) for f in filenames if f.endswith(".v4_auto_conll") \
                            or f.endswith(".v4_gold_conll")]
                cleaned_file_list = []
                for f in file_list:
                    fn = f.split('.')
                    if fn[1] == "v4_auto_conll":
                        gold = fn[0] + "." + "v4_gold_conll"
                        if gold not in file_list:
                            cleaned_file_list.append(f)
                    else:
                        cleaned_file_list.append(f)
            #self.load_file(file_list[0])
                doc_list = parallel_process(cleaned_file_list, read_file)
                for doc in doc_list:
                    kf.write(doc)

    def list_undetected_mentions(self, data_path, save_file, debug=True):
        self.read_corpus(data_path)
        print("🌋 Listing undetected mentions")
        with io.open(save_file, 'w', encoding='utf-8') as out_file:
            for doc in tqdm(self.docs):
                for name, part, utt_i, utt, coref in doc.missed_gold:
                    out_str = name + u"\t" + part + u"\t" + utt_i + u'\t"' + utt + u'"\n'
                    out_str += coref + u"\n"
                    out_file.write(out_str)
                    if debug: print(out_str)

    def read_corpus(self, data_path, debug=False):
        print("🌋 Reading files")
        for dirpath, _, filenames in os.walk(data_path):
            print("In", dirpath, os.path.abspath(dirpath))
            file_list = [os.path.join(dirpath, f) for f in filenames if f.endswith(".v4_auto_conll") \
                        or f.endswith(".v4_gold_conll")]
            cleaned_file_list = []
            for f in file_list:
                fn = f.split('.')
                if fn[1] == "v4_auto_conll":
                    gold = fn[0] + "." + "v4_gold_conll"
                    if gold not in file_list:
                        cleaned_file_list.append(f)
                else:
                    cleaned_file_list.append(f)
            doc_list = parallel_process(cleaned_file_list, load_file)
            for docs in doc_list:#executor.map(self.load_file, cleaned_file_list):
                for utts_text, utt_tokens, utts_corefs, utts_speakers, name, part in docs:
                    print("Imported", name)
                    if debug:
                        print("utts_text", utts_text)
                        print("utt_tokens", utt_tokens)
                        print("utts_corefs", utts_corefs)
                        print("utts_speakers", utts_speakers)
                        print("name, part", name, part)
                    self.utts_text += utts_text
                    self.utts_tokens += utt_tokens
                    self.utts_corefs += utts_corefs
                    self.utts_speakers += utts_speakers
                    self.utts_doc_idx += [len(self.docs_names)] * len(utts_text)
                    self.docs_names.append((name, part))
        print("utts_text size", len(self.utts_text))
        print("utts_tokens size", len(self.utts_tokens))
        print("utts_corefs size", len(self.utts_corefs))
        print("utts_speakers size", len(self.utts_speakers))
        print("utts_doc_idx size", len(self.utts_doc_idx))
        print("🌋 Building docs")
        for name, part in self.docs_names:
            self.docs.append(ConllDoc(name=name, part=part, nlp=None,
                                      blacklist=False, consider_speakers=True,
                                      embedding_extractor=self.embed_extractor,
                                      conll=CONLL_GENRES[name[:2]]))
        print("🌋 Loading spacy model")
        try:
            spacy.info('en_core_web_sm')
            model = 'en_core_web_sm'
        except IOError:
            print("No spacy 2 model detected, using spacy1 'en' model")
            spacy.info('en')
            model = 'en'
        nlp = spacy.load(model)
        print("🌋 Parsing utterances and filling docs")
        doc_iter = (s for s in self.utts_text)
        for utt_tuple in tqdm(zip(nlp.pipe(doc_iter),
                                           self.utts_tokens, self.utts_corefs,
                                           self.utts_speakers, self.utts_doc_idx)):
            spacy_tokens, conll_tokens, corefs, speaker, doc_id = utt_tuple
            if debug: print(unicode_(self.docs_names[doc_id]), "-", spacy_tokens)
            doc = spacy_tokens
            if debug: 
                out_str = "utterance " + unicode_(doc) + " corefs " + unicode_(corefs) + \
                          " speaker " + unicode_(speaker) + "doc_id" + unicode_(doc_id)
                print(out_str.encode('utf-8'))
            self.docs[doc_id].add_conll_utterance(doc, conll_tokens, corefs, speaker,
                                                  use_gold_mentions=self.use_gold_mentions)

    def build_and_gather_multiple_arrays(self, save_path):
        print("🌋 Extracting mentions features")
        parallel_process(self.docs, set_feats, n_jobs=self.n_jobs)

        print("🌋 Building and gathering arrays")
        arr =[{'doc': doc,
               'i': i} for i, doc in enumerate(self.docs)]
        arrays_dicts = parallel_process(arr, get_feats, use_kwargs=True, n_jobs=self.n_jobs)
        gathering_dict = dict((feat, None) for feat in FEATURES_NAMES)
        n_mentions_list = []
        pairs_ant_index = 0
        pairs_start_index = 0
        for n, p, arrays_dict in tqdm(arrays_dicts):
            for f in FEATURES_NAMES:
                if gathering_dict[f] is None:
                    gathering_dict[f] = arrays_dict[f]
                else:
                    if f == FEATURES_NAMES[6]:
                        array = [a + pairs_ant_index for a in arrays_dict[f]]
                    elif f == FEATURES_NAMES[3]:
                        array = [a + pairs_start_index for a in arrays_dict[f]]
                    else:
                        array = arrays_dict[f]
                    gathering_dict[f] += array
            pairs_ant_index += n
            pairs_start_index += p
            n_mentions_list.append(n)

        for feature in FEATURES_NAMES[:9]:
            print("Building numpy array for", feature, "length", len(gathering_dict[feature]))
            if feature != "mentions_spans":
                array = np.array(gathering_dict[feature])
                if array.ndim == 1:
                    array = np.expand_dims(array, axis=1)
            else:
                array = np.stack(gathering_dict[feature])
            # check_numpy_array(feature, array, n_mentions_list)
            print("Saving numpy", feature, "size", array.shape)
            np.save(save_path + feature, array)
        for feature in FEATURES_NAMES[9:]:
            print("Saving pickle", feature, "size", len(gathering_dict[feature]))
            with open(save_path + feature + '.bin', "wb") as fp:  
                pickle.dump(gathering_dict[feature], fp)

    def save_vocabulary(self, save_path, debug=False):
        def _vocabulary_to_file(path, vocabulary):
            print("🌋 Saving vocabulary")
            with io.open(path, "w", encoding='utf-8') as f:
                if debug:
                    print('voc saved in {path}, length: {length}'
                          .format(path=path, size=len(vocabulary)))
                for w in tunable_voc:
                    f.write(w + '\n')

        print("🌋 Building tunable vocabulary matrix from static vocabulary")
        tunable_voc = self.embed_extractor.tun_voc
        _vocabulary_to_file(
            path=save_path + 'tuned_word_vocabulary.txt',
            vocabulary=tunable_voc
        )

        static_voc = self.embed_extractor.stat_voc
        _vocabulary_to_file(
            path=save_path + 'static_word_vocabulary.txt',
            vocabulary=static_voc
        )

        tuned_word_embeddings = np.vstack([self.embed_extractor.get_stat_word(w)[1] for w in tunable_voc])
        print("Saving tunable voc, size:", tuned_word_embeddings.shape)
        np.save(save_path + "tuned_word_embeddings", tuned_word_embeddings)

        static_word_embeddings = np.vstack([self.embed_extractor.static_embeddings[w] for w in static_voc])
        print("Saving static voc, size:", static_word_embeddings.shape)
        np.save(save_path + "static_word_embeddings", static_word_embeddings)


if __name__ == '__main__':
    DIR_PATH = os.path.dirname(os.path.realpath(__file__))
    parser = argparse.ArgumentParser(description='Training the neural coreference model')
    parser.add_argument('--function', type=str, default='all', help='Function ("all", "key", "parse", "find_undetected")')
    parser.add_argument('--path', type=str, default=DIR_PATH + '/data/', help='Path to the dataset')
    parser.add_argument('--key', type=str, help='Path to an optional key file for scoring')
    parser.add_argument('--n_jobs', type=int, default=1, help='Number of parallel jobs (default 1)')
    args = parser.parse_args()
    if args.key is None:
        args.key = args.path + "/key.txt"
    CORPUS = ConllCorpus(n_jobs=args.n_jobs)
    if args.function == 'parse' or args.function == 'all':
        SAVE_DIR = args.path + "/numpy/"
        if not os.path.exists(SAVE_DIR):
            os.makedirs(SAVE_DIR)
        else:
            if os.listdir(SAVE_DIR):
                print("There are already data in", SAVE_DIR)
                print("Erasing")
                for file in os.listdir(SAVE_DIR):
                    print(file)
                    os.remove(SAVE_DIR + file)
        start_time = time.time()
        CORPUS.read_corpus(args.path)
        print('=> read_corpus time elapsed', time.time() - start_time)
        start_time2 = time.time()
        CORPUS.build_and_gather_multiple_arrays(SAVE_DIR)
        print('=> build_and_gather_multiple_arrays time elapsed', time.time() - start_time2)
        start_time2 = time.time()
        CORPUS.save_vocabulary(SAVE_DIR)
        print('=> save_vocabulary time elapsed', time.time() - start_time2)
        print('=> total time elapsed', time.time() - start_time)
    if args.function == 'key' or args.function == 'all':
        CORPUS.build_key_file(args.path, args.key)
    if args.function == 'find_undetected':
        CORPUS.list_undetected_mentions(args.path, args.path + "/undetected_mentions.txt")
