# Copyright 2022 Huawei Technologies Co., Ltd
# Copyright 2022 Aerospace Information Research Institute,
# Chinese Academy of Sciences.
#
# 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.
# ============================================================================
"""pretrain of ringmo"""
import os
import argparse
import aicc_tools as ac

from mindspore.train.model import Model

from ringmo_framework.lr import build_lr
from ringmo_framework.arch import build_model
from ringmo_framework.optim import build_optim
from ringmo_framework.datasets import build_dataset
from ringmo_framework.trainer import build_wrapper
from ringmo_framework.parallel_config import build_parallel_config
from ringmo_framework.tools.helper import count_params, build_context, str2bool
from ringmo_framework.tools.load_ckpt import load_ckpt
from ringmo_framework.monitors.callback import build_pretrain_callback
from register.config import RingMoConfig, ActionDict


@ac.aicc_monitor
def main(args):
    # init context
    cfts, profile_cb = build_context(args)

    # build dataset
    args.logger.info(".........Build Dataset..........")
    args.pretrain_dataset.data_path = cfts.get_dataset(args.pretrain_dataset.data_path)
    dataset = build_dataset(args)
    data_size = dataset.get_dataset_size()
    new_epochs = args.train_config.epoch
    if args.train_config.per_epoch_size and args.train_config.sink_mode:
        new_epochs = int((data_size / args.train_config.per_epoch_size) * new_epochs)
    else:
        args.train_config.per_epoch_size = data_size

    args.data_size = data_size
    args.logger.info("Will be Training epochs:{}, sink_size:{}".format(
        new_epochs, args.train_config.per_epoch_size))
    args.logger.info("Create training dataset finish, data size:{}".format(data_size))

    # build context config
    args.logger.info(".........Build context config..........")
    build_parallel_config(args)
    args.logger.info("context config is:{}".format(args.parallel_config))
    args.logger.info("moe config is:{}".format(args.moe_config))

    # build net
    args.logger.info(".........Build Net..........")
    net = build_model(args)
    args.logger.info("the params of the network: {} M.".format(count_params(net)))

    # build lr
    args.logger.info(".........Build LR Schedule..........")
    lr_schedule = build_lr(args)

    # define optimizer
    args.logger.info(".........Build Optimizer..........")
    optimizer = build_optim(args, net, lr_schedule, args.logger)

    # define model
    args.logger.info(".........Build Train Model..........")
    train_model = build_wrapper(args, net, optimizer, log=args.logger)
    args.logger.info("the params of the network: {} M.".format(count_params(train_model)))

    # define Model and begin training
    args.logger.info(".........Starting Init Train Model..........")
    model = Model(train_model)

    # resume ckpt
    load_ckpt(args, cfts, net, model, train_model, dataset, new_epochs)

    # define callback
    callback = build_pretrain_callback(args, cfts)

    if args.profile:
        callback.append(profile_cb)

    args.logger.info(".........Starting Training Model..........")
    model.train(new_epochs, dataset, callbacks=callback,
                dataset_sink_mode=args.train_config.sink_mode,
                sink_size=args.train_config.per_epoch_size)


if __name__ == "__main__":
    work_path = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        default=os.path.join(work_path, "config path"),
        help='YAML config files')
    parser.add_argument('--device_id', default=None, type=int, help='device id')
    parser.add_argument('--seed', default=None, type=int, help='random seed')
    parser.add_argument('--use_parallel', default=None, type=str2bool, help='whether use parallel mode')
    parser.add_argument('--profile', default=None, type=str2bool, help='whether use profile analysis')
    parser.add_argument(
        '--options',
        nargs='+',
        action=ActionDict,
        help='override some settings in the used config, the key-value pair'
             'in xxx=yyy format will be merged into config file')

    args_ = parser.parse_args()
    config = RingMoConfig(args_.config)
    if args_.device_id is not None:
        config.context.device_id = args_.device_id
    if args_.seed is not None:
        config.seed = args_.seed
    if args_.use_parallel is not None:
        config.use_parallel = args_.use_parallel
    if args_.profile is not None:
        config.profile = args_.profile
    if args_.options is not None:
        config.merge_from_dict(args_.options)

    main(config)