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()