import numpy as np
import pandas as pd
import torch
from torch.utils import data
import json
import collections
from sklearn.preprocessing import OneHotEncoder
from subword_nmt.apply_bpe import BPE
import codecs
from tqdm import tqdm
import csv


# vocab_path = './ESPF/protein_codes_uniprot.txt'
# bpe_codes_protein = codecs.open(vocab_path)
# pbpe = BPE(bpe_codes_protein, merges=-1, separator='')
# sub_csv = pd.read_csv('./ESPF/subword_units_map_uniprot.csv')
#
# idx2word_p = sub_csv['index'].values
# words2idx_p = dict(zip(idx2word_p, range(0, len(idx2word_p))))

# vocab_path = './ESPF/drug_codes_chembl.txt'
# bpe_codes_drug = codecs.open(vocab_path)
# dbpe = BPE(bpe_codes_drug, merges=-1, separator='')
# sub_csv = pd.read_csv('./ESPF/subword_units_map_chembl.csv')
#
# idx2word_d = sub_csv['index'].values
# words2idx_d = dict(zip(idx2word_d, range(0, len(idx2word_d))))

# max_d = 205
# max_p = 545


def load_vocab(vocab_file):
    """Loads a vocabulary file into a dictionary."""
    vocab = collections.OrderedDict()
    with open(vocab_file, "r", encoding="utf-8") as reader:
        tokens = reader.readlines()
    for index, token in enumerate(tokens):
        token = token.rstrip("\n")
        vocab[token] = index
    return vocab


def protein2emb_encoder(x, words2idx_p):
    max_p = 152
    # t1 = pbpe.process_line(x).split()  # split
    t1 = x.split(',')
    try:
        i1 = np.asarray([words2idx_p[i] for i in t1])  # index
    except:
        i1 = np.array([0])
        # print(x)

    l = len(i1)

    if l < max_p:
        i = np.pad(i1, (0, max_p - l), 'constant', constant_values=0)
        input_mask = ([1] * l) + ([0] * (max_p - l))
    else:
        i = i1[:max_p]
        input_mask = [1] * max_p

    return i, np.asarray(input_mask)


def drug2emb_encoder(x, dbpe, words2idx_d):
    max_d = 50
    # max_d = 100
    t1 = dbpe.process_line(x)
    t1 = t1.split()  # split
    try:
        i1 = np.asarray([words2idx_d[i] for i in t1])  # index
    except:
        i1 = np.array([0])
        # print(x)

    l = len(i1)
    print(i1)

    if l < max_d:
        i = np.pad(i1, (0, max_d - l), 'constant', constant_values=0)
        input_mask = ([1] * l) + ([0] * (max_d - l))

    else:
        i = i1[:max_d]
        input_mask = [1] * max_d

    return i, np.asarray(input_mask)


def seq2emb_encoder(input_seq, max_len, vocab):
    try:
        ids = np.asarray([vocab[i] for i in input_seq])
    except:
        ids = np.array([0])

    l = len(ids)

    if l < max_len:
        ids = np.pad(ids, (0, max_len - l), 'constant', constant_values=0)
        input_mask = np.array(([1] * l) + ([0] * (max_len - l)))
    else:
        ids = ids[:max_len]
        input_mask = np.array([1] * max_len)

    return ids, input_mask


class Data_Encoder(data.Dataset):
    def __init__(self, train_file, tokenizer_config):
        'Initialization'
        # load data
        with open(train_file["sps"], 'r') as f:
            self.sps = f.readlines()
        with open(train_file["smile"], 'r') as f:
            self.smile = f.readlines()
        with open(train_file["affinity"], 'r') as f:
            self.affinity = f.readlines()
        # define tokenizer
        self.begin_id = tokenizer_config["begin_id"]
        self.sep_id = tokenizer_config["separate_id"]
        self.max_len = tokenizer_config["max_len"]
        self.vocab = load_vocab(tokenizer_config["vocab_file"])
        bpe_codes_drug = codecs.open(tokenizer_config["vocab_pair"])
        self.dbpe = BPE(bpe_codes_drug, merges=-1, separator='')

    def __len__(self):
        'Denotes the total number of samples'
        return len(self.sps)

    def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        # Load data and get label
        # tokenization
        d = self.dbpe.process_line(self.smile[index].strip()).split()
        p = self.sps[index].strip().split(',')
        y = self.affinity[index].strip()

        input_seq = [self.begin_id] + d + [self.sep_id] + p + [self.sep_id]
        input, input_mask = seq2emb_encoder(input_seq, self.max_len, self.vocab)

        return input, y


if __name__ == "__main__":
    # local test
    # dataFolder = './IC50/SPS/train_smile'
    # with open(dataFolder, 'r') as f:
    #     train_smi = f.readlines()
    # drug_smi = train_smi[0]
    # d_v, input_mask_d = drug2emb_encoder(drug_smi)

    # test load vocab
    # vocab_file = './ESPF/vocab.txt'
    # vocab = load_vocab(vocab_file)

    # test train
    df_train = {"sps": './data/train_sps',
                "smile": './data/train_smile',
                "affinity": './data/train_ic50',
                }

    tokenizer_config = {"vocab_file": './config/vocab.txt',
                        "vocab_pair": './config/drug_codes_chembl.txt',
                        "begin_id": '[CLS]',
                        "separate_id": "[SEP]",
                        "max_len": 256
                        }
    params = {'batch_size': 5,
              'shuffle': True,
              'num_workers': 0,
              'drop_last': True}
    trainset = Data_Encoder(df_train, tokenizer_config)
    training_generator = data.DataLoader(trainset, **params)
    for i, (input, affinity) in tqdm(enumerate(training_generator)):
        print(input.size())
