import torch
import torch.onnx
from networks.mnist_LeNet import MNIST_LeNet_Autoencoder


def export_model_to_onnx(model, filepath, input_shape=(1, 1, 28, 28)):
    model.eval()  # Set the model to evaluation mode
    dummy_input = torch.randn(input_shape)  # Create a dummy input tensor
    torch.onnx.export(
        model,  # model being run
        dummy_input,  # model input (or a tuple for multiple inputs)
        filepath,  # where to save the model (can be a file or file-like object)
        export_params=True,  # store the trained parameter weights inside the model file
        opset_version=11,  # the ONNX version to export the model to
        do_constant_folding=True,  # whether to execute constant folding for optimization
        input_names=["input"],  # the model's input names
        output_names=["output"],  # the model's output names
        dynamic_axes={
            "input": {0: "batch_size"},  # variable length axes
            "output": {0: "batch_size"},
        },
    )


if __name__ == "__main__":
    # Initialize the autoencoder model
    autoencoder = MNIST_LeNet_Autoencoder(rep_dim=32)

    # Define the file path where the ONNX model will be saved
    onnx_file_path = "mnist_lenet_autoencoder.onnx"

    # Export the model
    export_model_to_onnx(autoencoder, onnx_file_path)

    print(f"Model has been exported to {onnx_file_path}")
