from transformers import AutoConfig, AutoModel, AutoTokenizer

import torch

from tqdm import tqdm

import numpy as np

import sys

sys.path.append("../../pretrain")

from load_umls import UMLS

from nltk.tokenize import word_tokenize

import ipdb

import os

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

log_list = 1 / np.log2(list(range(2, 1001, 1)))

batch_size = 512

max_seq_length = 32

# umls = UMLS("../../umls", source_range='SNOMEDCT_US')

t_list = ['Pharmacologic Substance', 'Disease or Syndrome',

          'Neoplastic Process', 'Clinical Drug', 'Finding', 'Injury or Poisoning']

def mcsm(embedding_list, embedding_type_list, type_list=t_list, k=40, lang_range=['ENG'], check_intersection=False):

    if check_intersection:

        if not os.path.exists("intersection.txt"):

            intersection_cui = get_intersection(

                embedding_list, embedding_type_list)

            with open("intersection.txt", "w", encoding="utf-8") as f:

                for cui in intersection_cui:

                    f.write(cui.strip() + "\n")

        else:

            with open("intersection.txt", "r", encoding="utf-8") as f:

                lines = f.readlines()

            intersection_cui = [line.strip() for line in lines]

    umls = UMLS("../../umls", source_range='SNOMEDCT_US',

                lang_range=lang_range)

    if check_intersection:

        cui_list = [cui for cui in intersection_cui

                    if cui in umls.cui2sty and umls.cui2sty[cui] in type_list]

    else:

        cui_list = [cui for cui, sty in umls.cui2sty.items()

                    if sty in type_list]

    opt = []

    for index, embedding in enumerate(embedding_list):

        if embedding_type_list[index].lower() == "cui":

            opt.append(mcsm_cui(embedding, umls, cui_list, type_list, k))

        if embedding_type_list[index].lower() == "word":

            opt.append(mcsm_word(embedding, umls, cui_list, type_list, k))

        if embedding_type_list[index].lower() == "bert":

            opt.append(mcsm_bert(embedding, umls, cui_list,

                                 type_list, k, summary_method="MEAN"))

            opt.append(mcsm_bert(embedding, umls, cui_list,

                                 type_list, k, summary_method="CLS"))

    return opt

def mcsm_cui(cui_embedding, umls, cui_list, type_list, k=40):

    w, _ = load_embedding(cui_embedding)

    if cui_list is None:

        cui_list = list(w.keys())

        print(f"Check cui count:{len(cui_list)}")

    else:

        print(f"All cui count:{len(cui_list)}")

        cui_list = list(set(w.keys()).intersection(set(cui_list)))

        print(f"Check cui count:{len(cui_list)}")

    term_embedding = np.array([w[cui] for cui in cui_list])

    term_type = [umls.cui2sty[cui] for cui in cui_list]

    return calculate_mcsm(term_embedding, term_type, type_list, k=k)

def mcsm_word(word_embedding, umls, cui_list, type_list, k=40):

    w, dim = load_embedding(word_embedding)

    print(f"All cui count:{len(cui_list)}")

    cui_str = [[word for word in word_tokenize(

        list(umls.cui2str[cui])[0]) if word in w] for cui in cui_list]

    check_count = 0

    term_type = []

    for index, cui in tqdm(enumerate(cui_str)):

        if len(cui) > 0:

            term_type.append(umls.cui2sty[cui_list[index]])

            tmp_emb = np.zeros((dim))

            for word in cui:

                tmp_emb += w[word]

            if check_count == 0:

                term_embedding = tmp_emb

            else:

                term_embedding = np.concatenate(

                    (term_embedding, tmp_emb), axis=0)

            check_count += 1

            """

            if check_count > 500:

                break

            """

    term_embedding = term_embedding.reshape((-1, dim))

    print(f"Check cui count:{check_count}")

    return calculate_mcsm(term_embedding, term_type, type_list, k=k)

def mcsm_bert(bert_embedding, umls, cui_list, type_list, k=40, summary_method="MEAN"):

    print(f"Check cui count:{len(cui_list)}")

    model, tokenizer = load_bert(bert_embedding)

    model.eval()

    input_ids = []

    for cui in tqdm(cui_list):

        input_ids.append(tokenizer.encode_plus(

            list(umls.cui2str[cui])[

                0], max_length=max_seq_length, add_special_tokens=True,

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

    count = len(input_ids)

    now_count = 0

    with tqdm(total=count) as pbar:

        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 = model(input_gpu_0)[1]

                if summary_method == "MEAN":

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

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

                if now_count == 0:

                    term_embedding = embed_np

                else:

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

                update = min(now_count + batch_size, count) - now_count

                now_count = now_count + update

                pbar.update(update)

    term_type = [umls.cui2sty[cui] for cui in cui_list]

    return calculate_mcsm(term_embedding, term_type, type_list, k=k)

def summary(opt):

    new_opt = {k: (np.mean(v), np.std(v)) for k, v in opt.items()}

    return new_opt

def calculate_mcsm(term_embedding, term_type, target_type_list, k):

    # term_embedding: term_count * embedding_dim

    # term_type: term_count

    term_embedding = torch.FloatTensor(term_embedding).to(device)

    embedding_norm = torch.norm(

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

    term_embedding = torch.div(term_embedding, embedding_norm)

    del embedding_norm

    output = {target_type: [] for target_type in target_type_list}

    for index, t in tqdm(enumerate(term_type)):

        if t in target_type_list:

            now = term_embedding[index]

            score = 0.0

            similarity = torch.matmul(term_embedding, now)

            # The most similar term is itself

            _, indices = torch.topk(similarity, k=k + 1)

            for i in range(1, k + 1, 1):

                if term_type[indices[i]] == t:

                    score += log_list[i - 1]

            output[t].append(score)

    del term_embedding

    output = summary(output)

    print(output)

    return output

def load_embedding(filename):

    print(filename)

    if filename.find('bin') >= 0:

        from gensim import models

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

        dim = W.vector_size

        return W, dim

    if filename.find('pkl') >= 0:

        import pickle

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

            W = pickle.load(f)

        for key, value in W.items():

            W[key] = np.array(list(map(float, value[1:-1].split(","))))

        dim = len(list(W.values())[0])

        return W, dim

    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(list(map(float, toks[1:])))

            W[w] = vec

    dim = len(list(W.values())[0])

    return W, dim

def load_bert(model_name_or_path):

    print(model_name_or_path)

    try:

        config = AutoConfig.from_pretrained(model_name_or_path)

        model = AutoModel.from_pretrained(

            model_name_or_path, config=config).to(device)

    except BaseException:

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

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

    try:

        model.output_hidden_states = False

    except BaseException:

        pass

    try:

        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

    except BaseException:

        tokenizer = AutoTokenizer.from_pretrained(

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

    return model, tokenizer

def get_intersection(embedding_list, embedding_type_list):

    intersection_cui = set()

    checker = True

    for index, embed in enumerate(embedding_list):

        if embedding_type_list[index] == "cui":

            w, _ = load_embedding(embed)

            if checker:

                intersection_cui = set(list(w.keys()))

                checker = False

            else:

                intersection_cui = set(

                    list(w.keys())).intersection(intersection_cui)

    print(f"Intersection count: {len(intersection_cui)}")

    return list(intersection_cui)

if __name__ == "__main__":

    """

    embedding_list = ["../../embeddings/claims_codes_hs_300.txt",

                      "../../embeddings/GoogleNews-vectors-negative300.bin",

                      "../../models/2020_eng"]

    #embedding_type_list = ["cui", "word", "bert"]

    embedding_list = ["../../embeddings/wikipedia-pubmed-and-PMC-w2v.bin",

                      "../../embeddings/bio_nlp_vec/PubMed-shuffle-win-2.bin",

                      "../../embeddings/bio_nlp_vec/PubMed-shuffle-win-30.bin"]

    embedding_type_list = ["word", "word", "word"]

    embedding_list = ["../../embeddings/DeVine_etal_200.txt",

                      "/home/yz/pretraining_models/cui2vec.pkl"]

    embedding_type_list = ["cui", "cui"]

    """

    #mcsm([embedding_list[2], embedding_type_list[2]])

    """

    embedding_list = ["../../embeddings/claims_codes_hs_300.txt",

                      "../../embeddings/DeVine_etal_200.txt",

                      "/home/yz/pretraining_models/cui2vec.pkl"]

    embedding_type_list = ["cui", "cui", "cui"]

    mcsm(embedding_list, embedding_type_list, check_intersection=True)

    """

    #embedding_list = ["../../models/2020_eng", "../../models/2020_all"]

    #mcsm(embedding_list, ["bert"] * 2, check_intersection=True)

    """

    embedding_list = ["../../embeddings/wikipedia-pubmed-and-PMC-w2v.bin",

                      "../../embeddings/GoogleNews-vectors-negative300.bin",

                      "../../embeddings/bio_nlp_vec/PubMed-shuffle-win-2.bin",

                      "../../embeddings/bio_nlp_vec/PubMed-shuffle-win-30.bin"]

    mcsm(embedding_list, ["word"] * 4, check_intersection=True)

    """

    embedding_list = ["/home/yz/pretraining_models/bert-base-cased",

                      "/home/yz/pretraining_models/biobert_v1.1",

                      "/home/yz/pretraining_models/BiomedNLP-PubMedBERT-base-uncased-abstract",

                      "/home/yz/pretraining_models/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext",

                      "/home/yz/pretraining_models/kexinghuang_clinical",

                      "emilyalsentzer/Bio_ClinicalBERT",

                      "../../models/UMLSBert_nosty"]

    #mcsm(embedding_list, ["bert"] * 6, check_intersection=True)

    #mcsm(embedding_list, ["bert"] * 6)

    mcsm([embedding_list[-1]], ["bert"])