from gensim import models
import os
import sys
sys.path.append("../../")
from pretrain.load_umls import UMLS
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModel, AutoConfig
from data_util import load
import tqdm

batch_size = 128
device = "cuda:0"


def get_umls():
    umls_label = []
    umls_label_set = set()
    umls_des = []
    umls = UMLS("../../umls", source_range=["MSH", "SNOMEDCT_US", "MDR"], only_load_dict=True)
    for cui in tqdm.tqdm(umls.cui2str):
        if not cui in umls_label_set:
            tmp_str = list(umls.cui2str[cui])
            umls_label.extend([cui] * len(tmp_str))
            umls_des.extend(tmp_str)
            umls_label_set.update([cui])
    print(len(umls_des))
    return umls_label, umls_des


def main(filename, summary_method, umls_label, umls_des):
    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, "../"))

    corpus_list = [("Medline", "es"), ("Medline", "fr"), ("Medline", "nl"), ("Medline", "de"),
                   ("EMEA", "es"), ("EMEA", "fr"), ("EMEA", "nl"), ("EMEA", "de"),
                   ("Patent", "fr"), ("Patent", "de")]
    """
    sty_list = ["Geographic Area",
                "Drug Delivery Device", "Medical Device", "Research Device",
                "Anatomical Abnormality", "Anatomical Structure", "Fully Formed Anatomical Structure",
                "Chemical", "Chemical Viewed Functionally", "Chemical Viewed Structurally", "Inorganic Chemical", "Organic Chemical", "Clinical Drug"]
    """
    result_dict = {}
    umls_embedding = get_bert_embed(umls_des, model, tokenizer, summary_method=summary_method, tqdm_bar=True)

    for corpus in corpus_list:
        output_text, output_label, label_set = load(dataset=corpus[0], lang=corpus[1])
        not_umls_label = [label for label in label_set if not label in umls_label]
        print(f"Count of not appearing in UMLS subset: {len(not_umls_label)}")
        text_embedding = get_bert_embed(output_text, model, tokenizer, summary_method=summary_method)
        predict_label = predict(text_embedding, umls_embedding, umls_label)
        p, r, f1 = metric(output_label, predict_label)
        result_dict[corpus[0] + "|" + corpus[1]] = (p, r, f1)
        print(p, r, f1)

    return result_dict

def predict(text_embedding, umls_embedding, umls_label):
    x_size = text_embedding.size(0)
    sim = torch.matmul(text_embedding, umls_embedding.t())
    most_similar = torch.max(sim, dim=1)[1]
    return [umls_label[idx] for idx in most_similar]


def metric(output_label, predict_label):
    predict_count = 0
    true_count = 0
    correct_count = 0
    for idx in range(len(output_label)):
        if isinstance(predict_label[idx], str):
            predict_label[idx] = [predict_label[idx]]
        if isinstance(output_label[idx], str):
            output_label[idx] = [output_label[idx]]
        predict_count += len(predict_label[idx])
        true_count += len(output_label[idx])
        for pred in predict_label[idx]:
            if pred in output_label[idx]:
                correct_count += 1

    p = correct_count / predict_count
    r = correct_count / true_count
    if p == 0. or r == 0.:
        f1 = 0.
    else:
        f1 = 2 * p * r / (p + r)
    return p, r, f1


def get_bert_embed(phrase_list, m, tok, normalize=True, summary_method="CLS", tqdm_bar=False):
    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():
        if tqdm_bar:
            pbar = tqdm.tqdm(total=count)
        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
            if now_count == 0:
                output = embed
            else:
                output = torch.cat((output, embed), dim=0)
            if tqdm_bar:
                pbar.update(min(now_count + batch_size, count) - now_count)
            now_count = min(now_count + batch_size, count)
        if tqdm_bar:
            pbar.close()
    return output


if __name__ == '__main__':
    umls_label, umls_des = get_umls()
    main("bert-base-multilingual-cased", "CLS", umls_label, umls_des)