# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved
#
# 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.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import ast
import argparse
import functools

import reader
from model_search import Network
from paddleslim.nas.darts import DARTSearch
sys.path[0] = os.path.join(os.path.dirname("__file__"), os.path.pardir)
from utility import add_arguments, print_arguments

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)

# yapf: disable
add_arg('log_freq',          int,   50,              "Log frequency.")
add_arg('use_multiprocess',  bool,  True,            "Whether use multiprocess reader.")
add_arg('batch_size',        int,   64,              "Minibatch size.")
add_arg('learning_rate',     float, 0.025,            "The start learning rate.")
add_arg('momentum',          float, 0.9,             "Momentum.")
add_arg('use_gpu',           bool,  True,            "Whether use GPU.")
add_arg('epochs',            int,   50,              "Epoch number.")
add_arg('init_channels',     int,   16,              "Init channel number.")
add_arg('layers',            int,   8,               "Total number of layers.")
add_arg('class_num',         int,   10,              "Class number of dataset.")
add_arg('trainset_num',      int,   50000,           "images number of trainset.")
add_arg('model_save_dir',    str,   'search_cifar', "The path to save model.")
add_arg('grad_clip',         float, 5,               "Gradient clipping.")
add_arg('arch_learning_rate',float, 3e-4,            "Learning rate for arch encoding.")
add_arg('method',            str,   'DARTS',         "The search method you would like to use")
add_arg('epochs_no_archopt', int,   0,               "Epochs not optimize the arch params")
add_arg('cutout_length',     int,   16,              "Cutout length.")
add_arg('cutout',            ast.literal_eval,  False, "Whether use cutout.")
add_arg('unrolled',          ast.literal_eval,  False, "Use one-step unrolled validation loss")
add_arg('use_data_parallel', ast.literal_eval,  False, "The flag indicating whether to use data parallel mode to train the model.")
# yapf: enable


def main(args):
    if not args.use_gpu:
        place = paddle.CPUPlace()
    elif not args.use_data_parallel:
        place = paddle.CUDAPlace(0)
    else:
        place = paddle.CUDAPlace(paddle.distributed.parallel.ParallelEnv()
                                 .dev_id)

    train_reader, valid_reader = reader.train_search(
        batch_size=args.batch_size,
        train_portion=0.5,
        is_shuffle=True,
        args=args)

    model = Network(args.init_channels, args.class_num, args.layers,
                    args.method)
    searcher = DARTSearch(
        model,
        train_reader,
        valid_reader,
        place,
        learning_rate=args.learning_rate,
        batchsize=args.batch_size,
        num_imgs=args.trainset_num,
        arch_learning_rate=args.arch_learning_rate,
        unrolled=args.unrolled,
        num_epochs=args.epochs,
        epochs_no_archopt=args.epochs_no_archopt,
        use_multiprocess=args.use_multiprocess,
        use_data_parallel=args.use_data_parallel,
        save_dir=args.model_save_dir,
        log_freq=args.log_freq)
    searcher.train()


if __name__ == '__main__':
    args = parser.parse_args()
    print_arguments(args)

    main(args)
