"""Command-line interface for caffe2onnx."""
import os
import argparse

from src.load_save_model import loadcaffemodel, saveonnxmodel
from src.caffe2onnx import Caffe2Onnx


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('caffe_graph_path', help="caffe's prototxt file path",
                        nargs='?', default="./caffemodel/test/test.prototxt")
    parser.add_argument('caffe_params_path', help="caffe's caffemodel file path",
                        nargs='?', default="./caffemodel/test/test.caffemodel")
    parser.add_argument('onnx_name', help="onnx model name",
                        nargs='?', default="test")
    parser.add_argument('save_dir', help="onnx model file saved path",
                        nargs='?', default="./onnxmodel/")
    return parser.parse_args()


def main():
    args = parse_args()
    caffe_graph_path = args.caffe_graph_path
    caffe_params_path = args.caffe_params_path
    onnx_name = args.onnx_name
    save_dir = args.save_dir
    save_path = os.path.join(save_dir, onnx_name + '.onnx')
    os.makedirs(save_dir, exist_ok=True)

    graph, params = loadcaffemodel(caffe_graph_path, caffe_params_path)
    c2o = Caffe2Onnx(graph, params, onnx_name)
    onnxmodel = c2o.createOnnxModel()
    saveonnxmodel(onnxmodel, save_path)


if __name__ == '__main__':
    main()
