import argparse
import pathlib

import torch
from omegaconf import OmegaConf

from app import create_app

if __name__ == "__main__":
    pathlib.Path("./checkpoints").mkdir(parents=True, exist_ok=True)
    print("===[Start]===")
    parser = argparse.ArgumentParser("Covid-EMR training script", add_help=False)
    parser.add_argument(
        "--cfg", type=str, required=True, metavar="FILE", help="path to config file"
    )
    parser.add_argument(
        "--cuda",
        type=int,
        required=False,
        metavar="CUDA NUMBER",
        help="gpu to train",
    )
    parser.add_argument(
        "--db",
        action="store_true",
        help="whether to connect database",
    )

    parser.add_argument(
        "--train",
        action="store_true",
        help="whether to train model, only execute inference stage if not",
    )

    args = parser.parse_args()
    print(f"===[{args.cfg}]===")
    conf = OmegaConf.load(args.cfg)
    conf.db = args.db
    conf.train = args.train

    # train on cpu by default
    device = torch.device("cpu")
    if args.cuda is not None:
        device = torch.device(
            f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu"
        )

    create_app(conf, device)
    print("===[End]===")