import sys

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

from linear_model import LinearModel

from load_umls import UMLS

import numpy as np

import os

import shutil

import torch

from torch.utils.data import DataLoader, TensorDataset, Dataset

from transformers import AutoTokenizer, AdamW, get_linear_schedule_with_warmup, AutoConfig, AutoModel

from time import time

from tqdm import tqdm

import ipdb

# parameters

embedding = sys.argv[1]

embedding_type = sys.argv[2]

freeze_embedding = sys.argv[3]

device = sys.argv[4]

if freeze_embedding.lower() in ['t', 'true']:

    freeze_embedding = True

else:

    freeze_embedding = False

if device == "0":

    device = "cuda:0"

if device == "1":

    device = "cuda:1"

if embedding_type == 'bert':

    epoch_num = 50

    if freeze_embedding:

        batch_size = 512

        learning_rate = 1e-3

    else:

        batch_size = 96

        learning_rate = 2e-5

    max_seq_length = 32

    try:

        tokenizer = AutoTokenizer.from_pretrained(embedding)

    except BaseException:

        tokenizer = AutoTokenizer.from_pretrained(

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

else:

    epoch_num = 50

    batch_size = 512

    learning_rate = 1e-3

    max_seq_length = 16

def pad(l):

    if len(l) > max_seq_length:

        return l[0:max_seq_length]

    return l + [use_embedding_count - 1] * (max_seq_length - len(l))

# load train and test

cui_train_0 = []

cui_train_1 = []

rel_train = []

with open("./data/x_train.txt") as f:

    lines = f.readlines()

    for line in lines:

        line = line.strip().split("\t")

        cui_train_0.append(line[0])

        cui_train_1.append(line[1])

with open("./data/y_train.txt") as f:

    lines = f.readlines()

    for line in lines:

        rel_train.append(line.strip())

cui_test_0 = []

cui_test_1 = []

rel_test = []

with open("./data/x_test.txt") as f:

    lines = f.readlines()

    for line in lines:

        line = line.strip().split("\t")

        cui_test_0.append(line[0])

        cui_test_1.append(line[1])

with open("./data/y_test.txt") as f:

    lines = f.readlines()

    for line in lines:

        rel_test.append(line.strip())

# build rel2id

rel_set = set(rel_train + rel_test)

rel2id = {rel: index for index, rel in enumerate(list(rel_set))}

id2rel = {index: rel for rel, index in rel2id.items()}

cui_set = set(cui_train_0 + cui_train_1 + cui_test_0 + cui_test_1)

print('Count of differnt cui:', len(cui_set))

# Deal cui type embedding

if embedding_type != 'bert':

    if embedding.find('txt') >= 0:

        with open(embedding, "r", encoding="utf-8") as f:

            line = f.readline()

            count, dim = map(int, line.strip().split())

            lines = f.readlines()

if embedding_type == 'cui':

    # build cui2id and use_embedding

    if embedding.find('txt') >= 0:

        cui2id = {}

        use_embedding_count = 0

        emb_sum = np.zeros(shape=(dim), dtype=float)

        for line in lines:

            l = line.strip().split()

            cui = l[0]

            if embedding.find('stanford') >= 0:

                cui = 'C' + cui

            emb = np.array(list(map(float, l[1:])))

            emb_sum += emb

            if cui in cui_set:

                cui2id[cui] = use_embedding_count

                if use_embedding_count == 0:

                    use_embedding = emb

                else:

                    use_embedding = np.concatenate((use_embedding, emb), axis=0)

                use_embedding_count += 1

        emb_avg = emb_sum / use_embedding_count

        use_embedding = np.concatenate((use_embedding, emb_avg), axis=0)

        use_embedding_count += 1

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

        print('Embedding shape:', use_embedding.shape)

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

        import pickle

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

            W = pickle.load(f)

        cui2id = {}

        use_embedding_count = 0

        dim = len(list(W.values())[0][1:-1].split(','))

        emb_sum = np.zeros(shape=(dim), dtype=float)

        for cui in cui_set:

            if cui in W and not cui in cui2id:

                emb = np.array([float(num) for num in W[cui][1:-1].split(',')])

                #ipdb.set_trace()

                emb_sum += emb

                cui2id[cui] = use_embedding_count

                if use_embedding_count == 0:

                    use_embedding = emb

                else:

                    use_embedding = np.concatenate((use_embedding, emb), axis=0)

                use_embedding_count += 1

        emb_avg = emb_sum / use_embedding_count

        if 'empty' in W:

            emb_avg = np.array([float(num) for num in W['empty'][1:-1].split(',')])

        use_embedding = np.concatenate((use_embedding, emb_avg), axis=0)

        use_embedding_count += 1

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

        print('Embedding shape:', use_embedding.shape)

    # apply cui2id and rel2id

    train_input_0 = [cui2id.get(cui, use_embedding_count - 1)

                     for cui in cui_train_0]

    train_input_1 = [cui2id.get(cui, use_embedding_count - 1)

                     for cui in cui_train_1]

    train_y = [rel2id[rel] for rel in rel_train]

    test_input_0 = [cui2id.get(cui, use_embedding_count - 1)

                    for cui in cui_test_0]

    test_input_1 = [cui2id.get(cui, use_embedding_count - 1)

                    for cui in cui_test_1]

    test_y = [rel2id[rel] for rel in rel_test]

# Find standard term name

if not embedding_type == 'cui':

    umls = UMLS("../../umls", only_load_dict=True)

    cui2str = {}

    #ipdb.set_trace()

    for cui in cui_set:

        standard_term = umls.search(code=cui, max_number=1)

        if standard_term is not None:

            cui2str[cui] = standard_term[0]

        else:

            cui2str[cui] = cui

# Deal word type embedding

if embedding_type == 'word':

    # tokenize

    from nltk.tokenize import word_tokenize

    cui2tokenize = {}

    for cui in cui2str:

        cui2tokenize[cui] = word_tokenize(cui2str[cui])

    # build word2id and use_embedding

    word2id = {}

    use_embedding_count = 0

    if embedding.find('txt') >= 0:

        emb_sum = np.zeros(shape=(dim), dtype=float)

        for line in lines:

            l = line.strip().split()

            word = l[0]

            emb = np.array(list(map(float, l[1:])))

            emb_sum += emb

            word2id[word] = use_embedding_count

            if use_embedding_count == 0:

                use_embedding = emb

            else:

                use_embedding = np.concatenate((use_embedding, emb), axis=0)

            use_embedding_count += 1

        emb_avg = emb_sum / use_embedding_count

        use_embedding = np.concatenate((use_embedding, emb_avg), axis=0)

        use_embedding_count += 1

        emb_zero = np.zeros_like(emb_avg)

        use_embedding = np.concatenate((use_embedding, emb_zero), axis=0)

        use_embedding_count += 1

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

        print('Embedding shape:', use_embedding.shape)

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

        import gensim

        model = gensim.models.KeyedVectors.load_word2vec_format(embedding, binary=True)

        emb_sum = np.zeros(shape=(model.vector_size), dtype=float)

        for cui in cui2tokenize:

            for w in cui2tokenize[cui]:

                if w in model and not w in word2id:

                    emb = model[w]

                    emb_sum += emb

                    word2id[w] = use_embedding_count

                    if use_embedding_count == 0:

                        use_embedding = emb

                    else:

                        use_embedding = np.concatenate((use_embedding, emb), axis=0)

                    use_embedding_count += 1  

        emb_avg = emb_sum / use_embedding_count

        use_embedding = np.concatenate((use_embedding, emb_avg), axis=0)

        use_embedding_count += 1

        emb_zero = np.zeros_like(emb_avg)

        use_embedding = np.concatenate((use_embedding, emb_zero), axis=0)

        use_embedding_count += 1

        use_embedding = use_embedding.reshape((-1, model.vector_size))

        print('Original embedding count:', len(model.wv.vocab))

        print('Embedding shape:', use_embedding.shape)                  

    # apply word2id and rel2id

    train_input_0 = [[word2id[w] for w in cui2tokenize[cui] if w in word2id] for cui in cui_train_0]

    train_input_1 = [[word2id[w] for w in cui2tokenize[cui] if w in word2id] for cui in cui_train_1]

    train_y = [rel2id[rel] for rel in rel_train]

    test_input_0 = [[word2id[w] for w in cui2tokenize[cui] if w in word2id] for cui in cui_test_0]

    test_input_1 = [[word2id[w] for w in cui2tokenize[cui] if w in word2id] for cui in cui_test_1]

    test_y = [rel2id[rel] for rel in rel_test]

    # average and padding

    # deal with input length = 0, use average

    train_input_0 = [cui if cui else [use_embedding_count - 2] for cui in train_input_0] 

    train_input_1 = [cui if cui else [use_embedding_count - 2] for cui in train_input_1]

    test_input_0 = [cui if cui else [use_embedding_count - 2] for cui in test_input_0]

    test_input_1 = [cui if cui else [use_embedding_count - 2] for cui in test_input_1]

    # calculate length

    train_length_0 = [len(cui) for cui in train_input_0]

    train_length_1 = [len(cui) for cui in train_input_1]

    test_length_0 = [len(cui) for cui in test_input_0]

    test_length_1 = [len(cui) for cui in test_input_1]

    # padding

    train_input_0 = list(map(pad, train_input_0))

    train_input_1 = list(map(pad, train_input_1))

    test_input_0 = list(map(pad, test_input_0))

    test_input_1 = list(map(pad, test_input_1))

# Deal bert type embedding

if embedding_type == 'bert':

    train_input_0 = []

    train_input_1 = []

    test_input_0 = []

    test_input_1 = []

    cui2tokenize = {}

    for cui in cui2str:

        cui2tokenize[cui] = tokenizer.encode_plus(

            cui2str[cui], max_length=max_seq_length, add_special_tokens=True,

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

    train_input_0 = [cui2tokenize[cui] for cui in cui_train_0]

    train_input_1 = [cui2tokenize[cui] for cui in cui_train_1]

    test_input_0 = [cui2tokenize[cui] for cui in cui_test_0]

    test_input_1 = [cui2tokenize[cui] for cui in cui_test_1]

    train_y = [rel2id[rel] for rel in rel_train]

    test_y = [rel2id[rel] for rel in rel_test]

# Dataset and Dataloader

train_input_0 = torch.LongTensor(train_input_0)

train_input_1 = torch.LongTensor(train_input_1)

test_input_0 = torch.LongTensor(test_input_0)

test_input_1 = torch.LongTensor(test_input_1)

train_y = torch.LongTensor(train_y)

test_y = torch.LongTensor(test_y)

if embedding_type != 'word':

    train_dataset = TensorDataset(train_input_0, train_input_1, train_y)

    test_dataset = TensorDataset(test_input_0, test_input_1, test_y)

else:

    train_length_0 = torch.LongTensor(train_length_0)

    train_length_1 = torch.LongTensor(train_length_1)

    test_length_0 = torch.LongTensor(test_length_0)

    test_length_1 = torch.LongTensor(test_length_1)

    train_dataset = TensorDataset(train_input_0, train_input_1, train_length_0, train_length_1, train_y)

    test_dataset = TensorDataset(test_input_0, test_input_1, test_length_0, test_length_1, test_y)

train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=1)

test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=1)

# Prepare model and optimizier

# model

if embedding_type != 'bert':

    use_embedding = torch.FloatTensor(use_embedding)

    model = LinearModel(len(rel2id), embedding_type, use_embedding, freeze_embedding).to(device)

else:

    try:

        config = AutoConfig.from_pretrained(embedding)

        bert_model = AutoModel.from_pretrained(embedding, config=config).to(device)

    except BaseException:

        bert_model = torch.load(os.path.join(embedding, 'pytorch_model.bin')).to(device)

    model = LinearModel(len(rel2id), embedding_type, bert_model, freeze_embedding).to(device)

# optimizier

if embedding_type != 'bert':

    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

if embedding_type == "bert":

    no_decay = ["bias", "LayerNorm.weight"]

    optimizer_grouped_parameters = [

        {

            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],

            "weight_decay": 0.0,

        },

        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0},

    ]

    optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate, eps=1e-8)

    scheduler = get_linear_schedule_with_warmup(optimizer,

                                                num_warmup_steps=int(epoch_num * len(train_dataloader) * 0.1),

                                                num_training_steps=epoch_num * len(train_dataloader))

# Prepare eval function

from sklearn.metrics import accuracy_score, classification_report, f1_score

def eval(m, dataloader):

    y_pred = []

    y_true = []

    m.eval()

    with torch.no_grad():

        for batch in dataloader:

            x0 = batch[0].to(device)

            x1 = batch[1].to(device)

            if m.embedding_type == "word":

                l0 = batch[2].to(device)

                l1 = batch[3].to(device)

                r = batch[4]

            else:

                l0 = l1 = None

                r = batch[2]

            pred, loss = m(x0, x1, l0, l1)

            y_pred += torch.max(pred, dim=1)[1].detach().cpu().numpy().tolist()

            y_true += r.detach().cpu().numpy().tolist()

    acc = accuracy_score(y_true, y_pred) * 100

    #f1 = f1_score(y_true, y_pred) * 100

    report = classification_report(y_true, y_pred)

    return acc, report

# Train and eval

if not os.path.exists("./result/"):

    os.mkdir("./result/")

for epoch_index in range(epoch_num):

    model.train()

    epoch_loss = 0.

    time_now = time()

    for batch in tqdm(train_dataloader):

        optimizer.zero_grad()

        x0 = batch[0].to(device)

        x1 = batch[1].to(device)

        if model.embedding_type == "word":

            l0 = batch[2].to(device)

            l1 = batch[3].to(device)

            r = batch[4].to(device)

        else:

            l0 = l1 = None

            r = batch[2].to(device)

        pred, loss = model(x0, x1, l0, l1, r)

        loss.backward()

        optimizer.step()

        if model.embedding_type == "bert":

            scheduler.step()

        epoch_loss += loss.item()

    print(epoch_index + 1, round(time() - time_now, 1), epoch_loss)

    acc, report = eval(model, test_dataloader)

    print("Accuracy:", acc)

    #print(report)