import os
import shutil
import onnx
from onnx import numpy_helper

tests = []

for d in os.listdir():
    if d.startswith("test_expand_"):
        tests.append(d)

if not os.path.exists("models"):
    os.mkdir("models")

if not os.path.exists("data"):
    os.mkdir("data")

for test in sorted(tests):
    model_path = os.path.join(test, "model.onnx")
    new_model_path = os.path.join("models", test + ".onnx")

    # copy input and output files
    inputs = []
    outputs = []
    for d in sorted(os.listdir(os.path.join(test, "test_data_set_0"))):
        if d.startswith("input_"):
            inputs.append(os.path.join(test, "test_data_set_0", d))

        if d.startswith("output_"):
            outputs.append(os.path.join(test, "test_data_set_0", d))
    #print(inputs)
    #print(outputs)

    def copy_to_data(filename_list):
        if len(filename_list) == 1:
            src = filename_list[0]
            dst = src.split("/")[-1].split("_")[:-1]
            dst.append(test + ".pb")
            dst = "_".join(dst)
            dst = os.path.join("data", dst)
            shutil.copyfile(src, dst)
        else:
            for src in filename_list:
                dst = src.split("/")[-1].split("_")
                dst.insert(1, test)
                dst = "_".join(dst)
                dst = os.path.join("data", dst)
                shutil.copyfile(src, dst)

    if len(inputs) == 2:
        model = onnx.load(model_path)
        graph = model.graph

        new_shape = onnx.load_tensor(inputs[-1])
        new_shape_name = graph.input[1].name

        new_shape_node = onnx.helper.make_node("Constant", [], [new_shape_name], value=new_shape)
        graph.node.insert(0, new_shape_node)
        graph.input.pop(1)
        new_model = onnx.helper.make_model(graph)
        onnx.checker.check_model(new_model)
        onnx.save(new_model, new_model_path)

        copy_to_data(inputs[:-1])
        copy_to_data(outputs)
    else:
        shutil.copyfile(model_path, new_model_path)
        copy_to_data(inputs)
        copy_to_data(outputs)