import argparse
import json
import spacy

import numpy as np
from allennlp.commands.elmo import ElmoEmbedder
from pathos.threading import ThreadPool


def flatten_elmo(elmo_embed_vecs):
    for p in elmo_embed_vecs[()].keys():
        sent_vec_list = elmo_embed_vecs[()][p]
        cont_vec = sent_vec_list[0]
        for i in range(1, len(sent_vec_list)):
            cont_vec = np.hstack((cont_vec, sent_vec_list[i]))
        elmo_embed_vecs[()][p] = cont_vec
    for p in elmo_embed_vecs[()].keys():
        concat_vec = elmo_embed_vecs[()][p][0]
        concat_vec = np.hstack((concat_vec, elmo_embed_vecs[()][p][1]))
        concat_vec = np.hstack((concat_vec, elmo_embed_vecs[()][p][2]))
        elmo_embed_vecs[()][p] = concat_vec
    return elmo_embed_vecs

def mean_flattened_elmo(flattened_embed_vecs):
    for p in flattened_embed_vecs[()].keys():
        flattened_embed_vecs[()][p] = np.mean(flattened_embed_vecs[()][p], axis=0)
    return flattened_embed_vecs

def main():
    parser = argparse.ArgumentParser(description="Flatten ELMo vectors generated by allennlp in suitable form")
    parser.add_argument("-el", "--elmo_vecs", required=True, help="ELMo embed vectors generated by allennlp")
    parser.add_argument("-d", "--dataset_name", required=True, help="Name of dataset, to be used in output filename (e.g. by3test)")
    parser.add_argument("-o", "--outdir", required=True, help="Path to output dir")
    args = vars(parser.parse_args())
    embed_vecs_path = args["elmo_vecs"]
    dataset_name = args["dataset_name"]
    outdir_path = args["outdir"]

    embed_vecs = np.load(embed_vecs_path, allow_pickle=True)
    flattened_elmo_vecs = flatten_elmo(embed_vecs)
    np.save(outdir_path+"/"+dataset_name+"_elmo_concat_para", flattened_elmo_vecs, allow_pickle=True)
    mean_elmo_vecs = mean_flattened_elmo(flattened_elmo_vecs)
    np.save(outdir_path+"/"+dataset_name+"_elmo_mean_para", mean_elmo_vecs, allow_pickle=True)

if __name__ == '__main__':
    main()