#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

import paddle
from ppgan.apps import StyleGANv2FittingPredictor

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_image", type=str, help="path to source image")

    parser.add_argument("--need_align",
                        action="store_true",
                        help="whether to align input image")

    parser.add_argument("--start_lr",
                        type=float,
                        default=0.1,
                        help="learning rate at the begin of training")

    parser.add_argument("--final_lr",
                        type=float,
                        default=0.025,
                        help="learning rate at the end of training")

    parser.add_argument("--latent_level",
                        type=int,
                        nargs="+",
                        default=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
                        help="indices of latent code for training")

    parser.add_argument("--step",
                        type=int,
                        default=100,
                        help="optimize iterations")

    parser.add_argument("--mse_weight",
                        type=float,
                        default=1,
                        help="weight of the mse loss")

    parser.add_argument("--pre_latent",
                        type=str,
                        default=None,
                        help="path to pre-prepared latent codes")

    parser.add_argument("--output_path",
                        type=str,
                        default='output_dir',
                        help="path to output image dir")

    parser.add_argument("--weight_path",
                        type=str,
                        default=None,
                        help="path to model checkpoint path")

    parser.add_argument("--model_type",
                        type=str,
                        default=None,
                        help="type of model for loading pretrained model")

    parser.add_argument("--size",
                        type=int,
                        default=1024,
                        help="resolution of output image")

    parser.add_argument("--style_dim",
                        type=int,
                        default=512,
                        help="number of style dimension")

    parser.add_argument("--n_mlp",
                        type=int,
                        default=8,
                        help="number of mlp layer depth")

    parser.add_argument("--channel_multiplier",
                        type=int,
                        default=2,
                        help="number of channel multiplier")

    parser.add_argument("--cpu",
                        dest="cpu",
                        action="store_true",
                        help="cpu mode.")

    args = parser.parse_args()

    if args.cpu:
        paddle.set_device('cpu')

    predictor = StyleGANv2FittingPredictor(
        output_path=args.output_path,
        weight_path=args.weight_path,
        model_type=args.model_type,
        seed=None,
        size=args.size,
        style_dim=args.style_dim,
        n_mlp=args.n_mlp,
        channel_multiplier=args.channel_multiplier)
    predictor.run(args.input_image,
                  need_align=args.need_align,
                  start_lr=args.start_lr,
                  final_lr=args.final_lr,
                  latent_level=args.latent_level,
                  step=args.step,
                  mse_weight=args.mse_weight,
                  pre_latent=args.pre_latent)