import cv2
import yaml
import argparse
import os
from torch.utils.data import DataLoader

from datasets.gl3d_dataset import GL3DDataset
from trainer import Trainer
from trainer_single_norel import SingleTrainerNoRel
from trainer_single import SingleTrainer


if __name__ == '__main__':
    # add argument parser
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='./configs/config.yaml')
    parser.add_argument('--dataset_dir', type=str, default='/mnt/nvme2n1/hyz/data/GL3D')
    parser.add_argument('--data_split', type=str, default='comb')
    parser.add_argument('--is_training', type=bool, default=True)
    parser.add_argument('--job_name', type=str, default='')
    parser.add_argument('--gpu', type=str, default='0')
    parser.add_argument('--start_cnt', type=int, default=0)
    parser.add_argument('--stage', type=int, default=1)
    args = parser.parse_args()

    # load global config
    with open(args.config, 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    # setup dataloader
    dataset = GL3DDataset(args.dataset_dir, config['network'], args.data_split, is_training=args.is_training)
    data_loader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=4)

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu


    if args.stage == 1:
        trainer = SingleTrainerNoRel(config, f'cuda:0', data_loader, args.job_name, args.start_cnt)
    elif args.stage == 2:
        trainer = SingleTrainer(config, f'cuda:0', data_loader, args.job_name, args.start_cnt)
    elif args.stage == 3:
        trainer = Trainer(config, f'cuda:0', data_loader, args.job_name, args.start_cnt)
    else:
        raise NotImplementedError()
        
    trainer.train()