import os
import shutil

import onnx

base_path = "./conformance"
dst_path = "./collected_conformance"

def remove_optinoal_output(onnx_file):
    model = onnx.load(onnx_file)
    graph = model.graph
    nodes = graph.node

    # remove node optional outputs (mean & InvStdDev)
    node = nodes[0]
    node.output.remove(node.output[1])
    node.output.remove(node.output[1])

    # remove graph optional outputs (mean & InvStdDev)
    graph.output.remove(graph.output[1])
    graph.output.remove(graph.output[1])

    new_model = onnx.helper.make_model(graph, producer_name="backend-test modified by github.com/opencv/opencv_extra")
    new_model = onnx.shape_inference.infer_shapes(new_model)
    return new_model

for d in os.listdir(base_path):
    if "expanded" in d or not d.startswith("test"):
        continue

    onnx_src = os.path.join(base_path, d, "model.onnx")
    onnx_dst = os.path.join(dst_path, "models", "{}.onnx".format(d))
    new_onnx = remove_optinoal_output(onnx_src)
    onnx.save(new_onnx, onnx_dst)

    in0_src = os.path.join(base_path, d, "test_data_set_0", "input_0.pb")
    in0_dst = os.path.join(dst_path, "data", "input_{}_0.pb".format(d))
    shutil.copy2(in0_src, in0_dst)

    in1_src = os.path.join(base_path, d, "test_data_set_0", "input_1.pb")
    in1_dst = os.path.join(dst_path, "data", "input_{}_1.pb".format(d))
    shutil.copy2(in1_src, in1_dst)

    in2_src = os.path.join(base_path, d, "test_data_set_0", "input_2.pb")
    in2_dst = os.path.join(dst_path, "data", "input_{}_2.pb".format(d))
    shutil.copy2(in2_src, in2_dst)

    out0_src = os.path.join(base_path, d, "test_data_set_0", "output_0.pb")
    out0_dst = os.path.join(dst_path, "data", "output_{}.pb".format(d))
    shutil.copy2(out0_src, out0_dst)