# Copyright 2021 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================

"""train_criteo."""
import argparse
import json
import os

from mindspore import context, Tensor, ParameterTuple
from mindspore.communication.management import init, get_rank, get_group_size
from mindspore.context import ParallelMode
from mindspore.nn import TrainOneStepCell
from mindspore.nn.optim import Adam
from mindspore.train import Model
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
from mindspore.train.serialization import load_checkpoint, load_param_into_net

from src.config import train_config, encoder_kw, decoder_kw
from src.dataset import create_dataset, ASRDataset
from src.model import Jasper, NetWithLossClass
from src.eval_callback import SaveCallback
from src.lr_generator import get_lr

parser = argparse.ArgumentParser(description='Jasper training')
parser.add_argument('--pre_trained_model_path', type=str, default='', help='Pretrained checkpoint path')
parser.add_argument('--is_distributed', action="store_true", default=False, help='Distributed training')
# parser.add_argument('--bidirectional', action="store_false", default=True, help='Use bidirectional RNN')
parser.add_argument('--device_target', type=str, default="GPU", choices=("GPU", "CPU"),
                    help='Device target, support GPU and CPU, Default: GPU')
args = parser.parse_args()

if __name__ == '__main__':

    rank_id = 0
    group_size = 1
    config = train_config
    data_sink = (args.device_target != "CPU")
    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, save_graphs=False)
    if args.device_target == "GPU":
        context.set_context(enable_graph_kernel=True)
    if args.is_distributed:
        init()
        rank_id = get_rank()
        group_size = get_group_size()
        context.reset_auto_parallel_context()
        context.set_auto_parallel_context(device_num=get_group_size(), parallel_mode=ParallelMode.DATA_PARALLEL,
                                          gradients_mean=True)

    with open(config.DataConfig.labels_path) as label_file:
        labels = json.load(label_file)

    # ds_train = create_dataset(audio_conf=config.DataConfig.SpectConfig,
    #                           manifest_filepath=config.DataConfig.train_manifest,
    #                           labels=labels, normalize=True, train_mode=True,
    #                           batch_size=config.DataConfig.batch_size, rank=rank_id, group_size=group_size)
    dataset = ASRDataset(audio_conf=config.DataConfig.SpectConfig,
                        manifest_filepath=config.DataConfig.train_manifest, 
                        labels=labels, batch_size=config.DataConfig.batch_size)
    for data in dataset:
        print(data)