import argparse


class HifiTestOptions:
    def __init__(self):
        self.parser = argparse.ArgumentParser()
        self.initialized = False

    def initialize(self):
        self.parser.add_argument('--model_name', type=str, default='er8_bs1', help='er8_bs1')
        self.parser.add_argument('--input_size', type=int, default='256')
        self.parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')

        self.parser.add_argument('--src_img_path', type=str, default='data/source/elon-musk1.jpg')
        self.parser.add_argument('--video_path', type=str, default='data/source/demo.mp4')
        self.parser.add_argument('--video_to_1080p', action="store_true", help='change video resolution to 1080p')
        self.parser.add_argument('--mode', type=str, default='default', help='default merge')
        self.parser.add_argument('--align_method', type=str, default='68', help='face align method:68 5class')

        self.parser.add_argument('--use_gfpgan', action="store_true", help='use gfpgan for sr or not')
        self.parser.add_argument('--sr_weight', type=float, default=1.0)

        self.parser.add_argument('--use_color_trans', action="store_true", help='use color transfer or not')
        self.parser.add_argument('--color_trans_mode', type=str, default='rct', help='rct lct mkl idt sot')

        self.parser.add_argument('--output_dir', type=str, default='data/output')


    def parse(self, save=True):
        if not self.initialized:
            self.initialize()
        self.opt = self.parser.parse_args()

        str_ids = self.opt.gpu_ids.split(',')
        self.opt.gpu_ids = []
        for str_id in str_ids:
            id = int(str_id)
            if id >= 0:
                self.opt.gpu_ids.append(id)

        return self.opt