""" Config class for training """
import argparse
import os
from functools import partial
import torch


def get_parser(name):
    """ make default formatted parser """
    parser = argparse.ArgumentParser(name, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    # print default value always
    parser.add_argument = partial(parser.add_argument, help=' ')
    return parser


def parse_gpus(gpus):
    if gpus == 'all':
        return list(range(torch.cuda.device_count()))
    else:
        return [int(s) for s in gpus.split(',')]


class BaseConfig(argparse.Namespace):
    def print_params(self, prtf=print):
        prtf("")
        prtf("Parameters:")
        for attr, value in sorted(vars(self).items()):
            prtf("{}={}".format(attr.upper(), value))
        prtf("")

    def as_markdown(self):
        """ Return configs as markdown format """
        text = "|name|value|  \n|-|-|  \n"
        for attr, value in sorted(vars(self).items()):
            text += "|{}|{}|  \n".format(attr, value)

        return text


class TrainConfig(BaseConfig):
    def build_parser(self):
        parser = get_parser("Train config")
        parser.add_argument('--name', default="/home/ubuntu/zhaoqianfei/UNet_Mini/LACB_Net_A/log")
        parser.add_argument('--batch_size', type=int, default=1, help='batch size')
        parser.add_argument('--input_channels', type=int, default=1, help='input channels')
        parser.add_argument('--n_classes', type=int, default=1, help='number classes')
        parser.add_argument('--lr', type=float, default=0.0025, help='lr for weights')
        parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
        parser.add_argument('--weight_decay', type=float, default=3e-4, help='weight decay')
        parser.add_argument('--grad_clip', type=float, default=5.,
                            help='gradient clipping for weights')
        parser.add_argument('--print_freq', type=int, default=20, help='print frequency')
        parser.add_argument('--gpus', default='0', help='gpu device ids separated by comma. '
                                                        '`all` indicates use all gpus.')
        parser.add_argument('--epochs', type=int, default=400, help='# of training epochs')
        parser.add_argument('--init_channels', type=int, default=12)
        parser.add_argument('--seed', type=int, default=2, help='random seed')
        parser.add_argument('--workers', type=int, default=4, help='# of workers')
        parser.add_argument('--training_summary_dir', default=".../model/unet")
        parser.add_argument('--training_checkpoint_prefix', default=".../model/unet")
        parser.add_argument('--testing_checkpoint_name', default=".../model/unet_400.pt")
        parser.add_argument('--testing_output_dir', default=".../result")
        parser.add_argument('--root_dir', default=".../data_train_test/lalel")
        parser.add_argument('--validing_checkpoint_prefix', default=".../model")

        return parser

    def __init__(self):
        parser = self.build_parser()
        args = parser.parse_args()
        super().__init__(**vars(args))

        self.path = os.path.join('train', self.name)
        self.gpus = parse_gpus(self.gpus)