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
from torch.utils.tensorboard import SummaryWriter
import os

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': False,
                       'num_workers': args.workers,
                       'drop_last': True
                       }
    data_generator = DataLoader(dataset, **data_loder_para)
    # creat model
    print('------------------creat model---------------------------')
    config = BertConfig.from_pretrained(args.config)
    model = BertAffinityModel(config)
    print('model name : BertAffinity')
    if args.task == 'train':
        opt = torch.optim.Adam(model.parameters(), lr=args.lr)
        loss_fct = torch.nn.MSELoss()
        writer = SummaryWriter('./log/' + args.log)
        num_step = args.epochs * len(data_generator)

        # detect GPU
        if torch.cuda.is_available():
            model.cuda(args.device)
        print('begin training')
        for epoch in range(args.epochs):
            for i, (input, affinity) in enumerate(data_generator):
                # use cuda
                if torch.cuda.is_available():
                    input.long().cuda(args.device)
                    affinity.cuda(args.divice)
                # input model
                pred_affinity = model(input)
                loss = loss_fct(pred_affinity, affinity.unsqueeze(-1))
                writer.add_scalar('loss', loss, num_step)
                # Update gradient
                opt.zero_grad()
                loss.backward()
                opt.step()

                if (i % 100 == 0):
                    print('Training at Epoch ' + str(epoch + 1) + ' step ' + str(i) + ' with loss ' + str(
                        loss.cpu().detach().numpy()))
                # save
                if epoch > args.epoches/3:
                    save_path = './model/' + args.save + '/'
                    if not os.path.exists(save_path):
                        os.mkdir(save_path)
                    torch.save(model.state_dict(), save_path + 'epoch-{}-step-{}-loss-{}.pth'.format(epoch, i, loss))







if __name__ == '__main__':
    # get parameter
    parser = ArgumentParser(description='BertAffinity')
    parser.add_argument('-b', '--batch-size', default=8, 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')
    parser.add_argument('--log', default='training_log', type=str, help='training log')
    parser.add_argument('--savedir', default='train', type=str, help='model save path')
    parser.add_argument('--device', default='0', type=str, help='name of GPU')
    args = parser.parse_args()


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

