"""
 Copyright (c) 2023, NVIDIA CORPORATION.
 
 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 hugectr2onnx.hugectr_loader import HugeCTRLoader, LayerParams
from hugectr2onnx.graph_builder import GraphBuilder
import argparse


def convert(
    onnx_model_path,
    graph_config,
    dense_model,
    convert_embedding=False,
    sparse_models=None,
    ntp_file=None,
    graph_name="hugectr",
):
    """Convert a HugeCTR model to an ONNX model
    Args:
        onnx_model_path: the path to store the ONNX model
        graph_config: the graph configuration JSON file of the HugeCTR model
        dense_model: the file of the dense weights for the HugeCTR model
        convert_embedding: whether to convert the sparse embeddings for the HugeCTR model (optional)
        sparse_models: the files of the sparse embeddings for the HugeCTR model (optional)
        ntp_file: the file of the non-trainable parameters for the HugeCTR model (optional)
        graph_name: the graph name for the ONNX model (optional)
    """
    loader = HugeCTRLoader(graph_config, dense_model, convert_embedding, sparse_models, ntp_file)
    builder = GraphBuilder(convert_embedding)
    for _ in range(loader.layers):
        layer_params, weights_dict, dimensions = loader.load_layer()
        print(f"[HUGECTR2ONNX][INFO]: Converting {layer_params.layer_type} layer to ONNX")
        builder.add_layer(layer_params, weights_dict, dimensions)
    builder.create_graph(graph_name)
    builder.save_model(onnx_model_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert HugeCTR model to ONNX.")
    parser.add_argument("--onnx_model_path", type=str, required=True, help="Output onnx model path")
    parser.add_argument(
        "--graph_config", type=str, required=True, help="HugeCTR graph configuration file"
    )
    parser.add_argument("--dense_model", type=str, required=True, help="HugeCTR dense model")
    parser.add_argument(
        "--convert_embedding",
        action="store_true",
        help="Convert sparse embedding or not (optional)",
    )
    parser.add_argument("--sparse_models", nargs="*", help="HugeCTR sparse models (optional)")
    parser.add_argument(
        "--ntp_file", type=str, default=None, help="HugeCTR non-trainable parameters (optional)"
    )
    parser.add_argument(
        "--graph_name", type=str, default="hugectr", help="Graph name for the ONNX model (optional)"
    )
    args = parser.parse_args()
    print(args)
    convert(
        args.onnx_model_path,
        args.graph_config,
        args.dense_model,
        args.convert_embedding,
        args.sparse_models,
        args.ntp_file,
        args.graph_name,
    )
