from argparse import ArgumentParser
from dataset import Data_Encoder
import torch
from torch.utils.data import DataLoader
from configuration_bert import BertConfig
from modeling_bert import BertAffinityModel


def get_task(task_name):
    if task_name.lower() == '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
                            }
        return df_train, tokenizer_config
    elif task_name.lower() == 'test':
        df_test = {"sps": './data/test_sps',
                    "smile": './data/test_smile',
                    "affinity": './data/test_ic50',
                    }
        tokenizer_config = {"vocab_file": './config/vocab.txt',
                            "vocab_pair": './config/drug_codes_chembl.txt',
                            "begin_id": '[CLS]',
                            "separate_id": "[SEP]",
                            "max_len": 256
                            }

        return df_test, tokenizer_config


def main(args):
    # load data
    data_file, tokenizer_config = get_task(args.task)
    dataset = Data_Encoder(data_file, tokenizer_config)
    data_loder_para = {'batch_size': args.batch_size,
                       'shuffle': True,
                       'num_workers': args.workers,
                       'drop_last': True
                       }
    data_generator = DataLoader(dataset, **data_loder_para)
    # creat model
    config = BertConfig.from_pretrained(args.config)
    model = BertAffinityModel(config)

    if args.task == 'train':
        opt = torch.optim.Adam(model.parameters(), lr=args.lr)
        loss_fct = torch.nn.MSELoss()
        for epoch in range(args.epochs):
            for i, (input, affinity) in enumerate(data_generator):
                pred_affinity = model(input.long())
                loss = loss_fct(pred_affinity, affinity)
                opt.zero_grad()
                loss.backward()
                opt.step()
                print('------------------')






if __name__ == '__main__':
    # get parameter
    parser = ArgumentParser(description='BertAffinity')
    parser.add_argument('-b', '--batch-size', default=1, type=int,
                        metavar='N',
                        help='mini-batch size (default: 16), this is the total '
                             'batch size of all GPUs on the current node when '
                             'using Data Parallel or Distributed Data Parallel')
    parser.add_argument('-j', '--workers', default=0, type=int, metavar='N',
                        help='number of data loading workers (default: 0)')
    parser.add_argument('--epochs', default=50, type=int, metavar='N',
                        help='number of total epochs to run')
    parser.add_argument('--task', choices=['train', 'test', 'channel', 'ER', 'GPCR', 'kinase'],
                        default='train', type=str, metavar='TASK',
                        help='Task name. Could be train, test, channel, ER, GPCR, kinase.')
    parser.add_argument('--lr', '--learning-rate', default=1e-4, type=float,
                        metavar='LR', help='initial learning rate', dest='lr')
    parser.add_argument('--config', default='./config/config.json', type=str, help='model config file path')
    args = parser.parse_args()


    # input dict
    main(args)
    # load_config(args.config)

