from gensim import models

import os

import sys

import torch

import numpy as np

from transformers import AutoTokenizer, AutoModel, AutoConfig

batch_size = 64

device = "cuda:1"

def main():

    filename = sys.argv[1]

    print(filename)

    bert_like = False

    if filename[-3:] in ["vec", "txt"]:

        W = load_vectors(filename, dev=False)

    elif filename[-3:] == "bin":

        W = load_vectors_bin(filename)

    else:

        bert_like = True

        try:

            config = AutoConfig.from_pretrained(filename)

            model = AutoModel.from_pretrained(

                filename, config=config).to(device)

        except BaseException:

            model = torch.load(os.path.join(

                filename, 'pytorch_model.bin')).to(device)

        try:

            model.output_hidden_states = False

        except BaseException:

            pass

        try:

            tokenizer = AutoTokenizer.from_pretrained(filename)

        except BaseException:

            tokenizer = AutoTokenizer.from_pretrained(

                os.path.join(filename, "../"))

    top_k = 3

    if bert_like:

        eval(model, tokenizer, './cadec/data/cadec', top_k=top_k, summary_method="CLS")

        eval(model, tokenizer, './cadec/data/cadec', top_k=top_k, summary_method="MEAN")

        eval(model, tokenizer, './cadec/data/psytar_disjoint_folds', top_k=top_k, summary_method="CLS")

        eval(model, tokenizer, './cadec/data/psytar_disjoint_folds', top_k=top_k, summary_method="MEAN")

    else:

        eval(W, None, './cadec/data/cadec', top_k=top_k)

        eval(W, None, './cadec/data/psytar_disjoint_folds', top_k=top_k)

def get_bert_embed(phrase_list, m, tok, normalize=True, summary_method="CLS"):

    input_ids = []

    for phrase in phrase_list:

        input_ids.append(tok.encode_plus(

            phrase, max_length=32, add_special_tokens=True,

            truncation=True, pad_to_max_length=True)['input_ids'])

    m.eval()

    count = len(input_ids)

    now_count = 0

    with torch.no_grad():

        while now_count < count:

            input_gpu_0 = torch.LongTensor(input_ids[now_count:min(

                now_count + batch_size, count)]).to(device)

            if summary_method == "CLS":

                embed = m(input_gpu_0)[1]

            if summary_method == "MEAN":

                embed = torch.mean(m(input_gpu_0)[0], dim=1)

            if normalize:

                embed_norm = torch.norm(

                    embed, p=2, dim=1, keepdim=True).clamp(min=1e-12)

                embed = embed / embed_norm

            embed_np = embed.cpu().detach().numpy()

            if now_count == 0:

                output = embed_np

            else:

                output = np.concatenate((output, embed_np), axis=0)

            now_count = min(now_count + batch_size, count)

    return output

def eval_one(m, tok, folder, top_k, summary_method=None):

    with open(os.path.join(folder, "standard.txt"), "r", encoding="utf-8") as f:

        lines = f.readlines()

    label2id = {line.strip().split(

        "\t")[0]: index for index, line in enumerate(lines)}

    standard_lines = [line.strip().split("\t") for line in lines]

    #standard_feat = np.array([get_bert_embed(text, m, tok) for (label, text) in standard_lines])

    if tok is not None:

        standard_feat = get_bert_embed(

            [text for (label, text) in standard_lines], m, tok, normalize=True, summary_method=summary_method)

    else:

        standard_feat = embed(

            [text for (label, text) in standard_lines], m.vector_size, m)

    with open(os.path.join(folder, "test.txt"), "r", encoding="utf-8") as f:

        lines = f.readlines()

    test_lines = [line.strip().split("\t") for line in lines]

    #test_feat = np.array([get_bert_embed(text, m, tok) for (label, text) in test_lines])

    if tok is not None:

        test_feat = get_bert_embed(

            [text for (label, text) in test_lines], m, tok, normalize=True, summary_method=summary_method)

    else:

        test_feat = embed(

            [text for (label, text) in test_lines], m.vector_size, m)

    sim_mat = np.dot(test_feat, standard_feat.T)

    correct_1 = 0

    correct_k = 0

    pred_top_k = torch.topk(torch.FloatTensor(sim_mat), k=top_k)[

        1].cpu().numpy()

    for i in range(len(test_lines)):

        true_id = label2id[test_lines[i][0]]

        if pred_top_k[i][0] == true_id:

            correct_1 += 1

        if true_id in list(pred_top_k[i]):

            correct_k += 1

    acc_1 = correct_1 / len(test_lines)

    acc_k = correct_k / len(test_lines)

    return acc_1, acc_k

def eval(m, tok, task_name, top_k=3, summary_method=None):

    acc_1_list = []

    acc_k_list = []

    for p in os.listdir(task_name):

        acc_1, acc_k = eval_one(m, tok, os.path.join(task_name, p), top_k, summary_method=summary_method)

        acc_1_list.append(acc_1)

        acc_k_list.append(acc_k)

    print(task_name, summary_method)

    print(f"top_k={top_k}")

    print(acc_1_list)

    print(acc_k_list)

    print(sum(acc_1_list) / 5, sum(acc_k_list) / 5)

    return None

def load_vectors(filename):

    W = {}

    with open(filename, 'r') as f:

        for i, line in enumerate(f.readlines()):

            if i == 0:

                continue

            toks = line.strip().split()

            w = toks[0]

            vec = np.array(map(float, toks[1:]))

            W[w] = vec

    return W

def load_vectors_bin(filename):

    w = models.KeyedVectors.load_word2vec_format(filename, binary=True)

    return w

def cosine(u, v):

    return np.dot(u, v)

def norm(v):

    return np.dot(v, v)**0.5

def embed_one(phrase, dim, W):

    words = phrase.split()

    vectors = [W[w] for w in words if (w in W)]

    v = sum(vectors, np.zeros(dim))

    return v / (norm(v) + 1e-9)

def embed(phrase_list, dim, W):

    return np.array([embed_one(phrase, dim, W) for phrase in phrase_list])

if __name__ == '__main__':

    main()