"""
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: train.py
@time: 2020-05-27 08:52:31
@desc:
"""
import os
from jjzhk.config import ZKCFG
from ELib.slover import BelodonSolver
import argparse
import torch

if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')


def main(dt, mo, bb, dr, ph):
    if bb == "":
        cfgfile = os.path.join("%s" % dt, "%s.yml" % mo)
    else:
        cfgfile = os.path.join("%s" % dt, "%s_%s.yml" % (mo, backbone))
    cfg = ZKCFG(basefile="basic_%s.yml" % dt, cfgfile=cfgfile, rootpath="cfgs")

    cfg.BASE.DATAROOT = dr
    cfg.BASE.DATA_ROOT = os.path.join(cfg.BASE.DATAROOT, cfg.BASE.DATA_ROOT)

    solver = BelodonSolver(cfg)

    if ph == "train":
        solver.train()
    elif ph == "eval":
        solver.eval()
    elif ph == "test":
        solver.test()


if __name__ == '__main__':
    # main("voc", "fcn8satonce", "", "test")
    # main("voc", "deeplabv3", "xception65", "train")
    # 1, 21, 524, 932

    # model = "deeplabv3"  # args.model
    # datatype = "voc"  # args.datatype
    # backbone = "resnet50"  # args.net
    # dataroot = '/Users/JJZHK/data/'  # args.dataroot
    # phase = 'train'  # args.phase

    # model = "fcn32s"  # args.model
    # datatype = "cityscapes"  # args.datatype
    # backbone = ""  # args.net
    # dataroot = '/Users/JJZHK/data/'  # args.dataroot
    # phase = 'train'  # args.phase

    parser = argparse.ArgumentParser(
        description='Segmentation Project')
    parser.add_argument('-dataroot', default='/Users/JJZHK/data/', type=str, help='')
    parser.add_argument('-model', default='ssd', type=str, help='')
    parser.add_argument('-datatype', default='voc', type=str, help='')
    parser.add_argument('-net', default='vgg16', type=str, help='')
    parser.add_argument('-phase', default='train', type=str, help='')
    args = parser.parse_args()
    model = args.model
    datatype = args.datatype
    backbone = args.net
    dataroot = args.dataroot
    phase = args.phase

    main(datatype, model, backbone, dataroot, phase)