/*
 * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
 * Licensed under the MIT License.
 */
package ai.onnxruntime;

import ai.onnxruntime.OnnxMl.StringStringEntryProto;
import ai.onnxruntime.OnnxMl.TensorProto.DataLocation;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;

/** Methods to generate test models. */
public final class ModelGenerators {
  private ModelGenerators() {}

  public static OnnxMl.TensorShapeProto getShapeProto(
      long[] dimensions, String[] dimensionOverrides) {
    OnnxMl.TensorShapeProto.Builder builder = OnnxMl.TensorShapeProto.newBuilder();
    for (int i = 0; i < dimensions.length; i++) {
      if (dimensions[i] == -1) {
        builder.addDim(
            OnnxMl.TensorShapeProto.Dimension.newBuilder()
                .setDimParam(dimensionOverrides[i])
                .build());
      } else {
        builder.addDim(
            OnnxMl.TensorShapeProto.Dimension.newBuilder().setDimValue(dimensions[i]).build());
      }
    }
    return builder.build();
  }

  public static OnnxMl.TypeProto buildTensorTypeNode(
      long[] dimensions, String[] dimensionOverrides, OnnxMl.TensorProto.DataType type) {
    OnnxMl.TypeProto.Builder builder = OnnxMl.TypeProto.newBuilder();

    OnnxMl.TypeProto.Tensor.Builder tensorBuilder = OnnxMl.TypeProto.Tensor.newBuilder();
    tensorBuilder.setElemType(type.getNumber());
    tensorBuilder.setShape(getShapeProto(dimensions, dimensionOverrides));
    builder.setTensorType(tensorBuilder.build());

    return builder.build();
  }

  public void generateExternalMatMul() throws IOException {
    OnnxMl.GraphProto.Builder graph = OnnxMl.GraphProto.newBuilder();
    graph.setName("ort-test-matmul");

    // Add placeholders
    OnnxMl.ValueInfoProto.Builder input = OnnxMl.ValueInfoProto.newBuilder();
    input.setName("input");
    OnnxMl.TypeProto inputType =
        buildTensorTypeNode(
            new long[] {-1, 4},
            new String[] {"batch_size", null},
            OnnxMl.TensorProto.DataType.FLOAT);
    input.setType(inputType);
    graph.addInput(input);
    OnnxMl.ValueInfoProto.Builder output = OnnxMl.ValueInfoProto.newBuilder();
    output.setName("output");
    OnnxMl.TypeProto outputType =
        buildTensorTypeNode(
            new long[] {-1, 4},
            new String[] {"batch_size", null},
            OnnxMl.TensorProto.DataType.FLOAT);
    output.setType(outputType);
    graph.addOutput(output);

    // Add initializers
    OnnxMl.TensorProto.Builder tensor = OnnxMl.TensorProto.newBuilder();
    tensor.addDims(4);
    tensor.addDims(4);
    tensor.setDataLocation(DataLocation.EXTERNAL);
    tensor.addExternalData(
        StringStringEntryProto.newBuilder()
            .setKey("location")
            .setValue("external-matmul.out")
            .build());
    tensor.addExternalData(
        StringStringEntryProto.newBuilder().setKey("offset").setValue("0").build());
    tensor.addExternalData(
        StringStringEntryProto.newBuilder().setKey("length").setValue("64").build());
    float[] floats =
        new float[] {1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f, 9f, 10f, 11f, 12f, 13f, 14f, 15f, 16f};
    ByteBuffer buf = ByteBuffer.allocate(64).order(ByteOrder.LITTLE_ENDIAN);
    FloatBuffer floatBuf = buf.asFloatBuffer();
    floatBuf.put(floats);
    floatBuf.rewind();
    buf.rewind();
    try (OutputStream os =
        Files.newOutputStream(Paths.get("src", "test", "resources", "external-matmul.out"))) {
      os.write(buf.array());
    }
    tensor.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
    tensor.setName("tensor");
    graph.addInitializer(tensor);

    // Add operations
    OnnxMl.NodeProto.Builder matmul = OnnxMl.NodeProto.newBuilder();
    matmul.setName("matmul-0");
    matmul.setOpType("MatMul");
    matmul.addInput("input");
    matmul.addInput("tensor");
    matmul.addOutput("output");
    graph.addNode(matmul);

    // Build model
    OnnxMl.ModelProto.Builder model = OnnxMl.ModelProto.newBuilder();
    model.setGraph(graph);
    model.setDocString("ORT matmul test");
    model.setModelVersion(0);
    model.setIrVersion(8);
    model.setDomain("ai.onnxruntime.test");
    model.addOpsetImport(OnnxMl.OperatorSetIdProto.newBuilder().setVersion(18).build());
    try (OutputStream os =
        Files.newOutputStream(Paths.get("src", "test", "resources", "java-external-matmul.onnx"))) {
      model.build().writeTo(os);
    }
  }

  public void generateMatMul() throws IOException {
    OnnxMl.GraphProto.Builder graph = OnnxMl.GraphProto.newBuilder();
    graph.setName("ort-test-matmul");

    // Add placeholders
    OnnxMl.ValueInfoProto.Builder input = OnnxMl.ValueInfoProto.newBuilder();
    input.setName("input");
    OnnxMl.TypeProto inputType =
        buildTensorTypeNode(
            new long[] {-1, 4},
            new String[] {"batch_size", null},
            OnnxMl.TensorProto.DataType.FLOAT);
    input.setType(inputType);
    graph.addInput(input);
    OnnxMl.ValueInfoProto.Builder output = OnnxMl.ValueInfoProto.newBuilder();
    output.setName("output");
    OnnxMl.TypeProto outputType =
        buildTensorTypeNode(
            new long[] {-1, 4},
            new String[] {"batch_size", null},
            OnnxMl.TensorProto.DataType.FLOAT);
    output.setType(outputType);
    graph.addOutput(output);

    // Add initializers
    OnnxMl.TensorProto.Builder tensor = OnnxMl.TensorProto.newBuilder();
    tensor.addDims(4);
    tensor.addDims(4);
    Float[] floats =
        new Float[] {1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f, 9f, 10f, 11f, 12f, 13f, 14f, 15f, 16f};
    tensor.addAllFloatData(Arrays.asList(floats));
    tensor.setDataType(OnnxMl.TensorProto.DataType.FLOAT.getNumber());
    tensor.setName("tensor");
    graph.addInitializer(tensor);

    // Add operations
    OnnxMl.NodeProto.Builder matmul = OnnxMl.NodeProto.newBuilder();
    matmul.setName("matmul-0");
    matmul.setOpType("MatMul");
    matmul.addInput("input");
    matmul.addInput("tensor");
    matmul.addOutput("output");
    graph.addNode(matmul);

    // Build model
    OnnxMl.ModelProto.Builder model = OnnxMl.ModelProto.newBuilder();
    model.setGraph(graph);
    model.setDocString("ORT matmul test");
    model.setModelVersion(0);
    model.setIrVersion(8);
    model.setDomain("ai.onnxruntime.test");
    model.addOpsetImport(OnnxMl.OperatorSetIdProto.newBuilder().setVersion(18).build());
    try (OutputStream os =
        Files.newOutputStream(Paths.get("src", "test", "resources", "java-matmul.onnx"))) {
      model.build().writeTo(os);
    }
  }
}
