from .base_options import BaseOptions


class TestOptions(BaseOptions):
    """This class includes test options.

    It also includes shared options defined in BaseOptions.
    """

    def initialize(self, parser):
        parser = BaseOptions.initialize(self, parser)  # define shared options
        parser.add_argument('--image_path', type=str, help='path to image for generate.')
        parser.add_argument('--ntest', type=int, default=float("inf"), help='# of test examples.')
        parser.add_argument('--aspect_ratio', type=float, default=1.0, help='aspect ratio of result images')
        parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
        parser.add_argument('--which_epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
        parser.add_argument('--how_many', type=int, default=50, help='how many test images to run')
        parser.add_argument('--results_dir', type=str, default='./results/', help='saves results here.')
        # Dropout and Batchnorm has different behavioir during training and test.
        parser.add_argument('--eval', action='store_true', help='use eval mode during test time.')
        parser.add_argument('--num_test', type=int, default=50, help='how many test images to run')
        # rewrite devalue values
        # To avoid cropping, the load_size should be the same as crop_size
        parser.set_defaults(load_size=parser.get_default('crop_size'))
        self.isTrain = False
        return parser



def generate():
    """
    generate single image specific by image path, and show the after generated image
    :return:
    """
    image_path = opt.image_path
    print('generate from {}'.format(image_path))

    data = load_image_for_prediction(opt, image_path)

    model = OneDirectionTestModel()
    model.initialize(opt=opt)
    model.set_input(data)
    model.test()

    visuals = model.get_current_visuals()
    generated_a = visuals['fake_B']

    image_generated = Image.fromarray(generated_a)
    image_generated.save(str(os.path.basename(image_path).split('.')[0]) + '_fake_b.jpg')

    combined_result = np.concatenate([img for _, img in visuals.items()], 1)
    image_combined = Image.fromarray(combined_result)
    image_combined.save(str(os.path.basename(image_path).split('.')[0]) + '_combined.jpg')
    image_combined.show()
    print('generated image saved.')



if __name__ == '__main__':
    generate()

def parse(self):
    if not self.initialized:
        self.initialize()
    self.opt = self.parser.parse_args()
    self.opt.isTrain = self.isTrain  # train or test

    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)

    args = vars(self.opt)

    print('------------ Options -------------')
    for k, v in sorted(args.items()):
        print('%s: %s' % (str(k), str(v)))
    print('-------------- End ----------------')

    # save to the disk
    expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
    util.mkdirs(expr_dir)
    file_name = os.path.join(expr_dir, 'opt.txt')
    with open(file_name, 'wt') as opt_file:
        opt_file.write('------------ Options -------------\n')
        for k, v in sorted(args.items()):
            opt_file.write('%s: %s\n' % (str(k), str(v)))
        opt_file.write('-------------- End ----------------\n')
    return self.opt