package com.tobe;


import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import io.grpc.StatusRuntimeException;
import org.tensorflow.framework.TensorProto;
import org.tensorflow.framework.TensorShapeProto;
import tensorflow.serving.Model;
import tensorflow.serving.Predict;
import tensorflow.serving.PredictionServiceGrpc;

import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * The general predit client for TensorFlow models.
 */
public class PredictClient {
    private static final Logger logger = Logger.getLogger(PredictClient.class.getName());
    private final ManagedChannel channel;
    private final PredictionServiceGrpc.PredictionServiceBlockingStub blockingStub;

    // Initialize gRPC client
    public PredictClient(String host, int port) {
        channel = ManagedChannelBuilder.forAddress(host, port)
                // Channels are secure by default (via SSL/TLS). For the example we disable TLS to avoid
                // needing certificates.
                .usePlaintext(true)
                .build();
        blockingStub = PredictionServiceGrpc.newBlockingStub(channel);
    }

    public static void main(String[] args) {
        System.out.println("Start the predict client");

        String host = "127.0.0.1";
        int port = 9000;
        String modelName = "cancer";
        long modelVersion = 1;

        // Parse command-line arguments
        if (args.length == 4) {
            host = args[0];
            port = Integer.parseInt(args[1]);
            modelName = args[2];
            modelVersion = Long.parseLong(args[3]);
        }

        // Run predict client to send request
        PredictClient client = new PredictClient(host, port);

        try {
            client.do_predict(modelName, modelVersion);
        } catch (Exception e) {
            System.out.println(e);
        } finally {
            try {
                client.shutdown();
            } catch (Exception e) {
                System.out.println(e);
            }
        }

        System.out.println("End of predict client");
    }

    public void shutdown() throws InterruptedException {
        channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
    }

    public void do_predict(String modelName, long modelVersion) {
        // Generate keys TensorProto
        int[][] keysTensorData = new int[][]{
                {1},
                {2}
        };

        TensorProto.Builder keysTensorBuilder = TensorProto.newBuilder();
        keysTensorBuilder.setDtype(org.tensorflow.framework.DataType.DT_INT32);

        for (int i = 0; i < keysTensorData.length; ++i) {
            for (int j = 0; j < keysTensorData[i].length; ++j) {
                keysTensorBuilder.addIntVal(keysTensorData[i][j]);
            }

        }
        TensorProto keysTensorProto = keysTensorBuilder.build();

        // Generate features TensorProto
        float[][] featuresTensorData = new float[][]{
                {10f, 10f, 10f, 8f, 6f, 1f, 8f, 9f, 1f},
                {10f, 10f, 10f, 8f, 6f, 1f, 8f, 9f, 1f},
        };

        TensorProto.Builder featuresTensorBuilder = TensorProto.newBuilder();
        featuresTensorBuilder.setDtype(org.tensorflow.framework.DataType.DT_FLOAT);

        for (int i = 0; i < featuresTensorData.length; ++i) {
            for (int j = 0; j < featuresTensorData[i].length; ++j) {
                featuresTensorBuilder.addFloatVal(featuresTensorData[i][j]);
            }
        }

        TensorShapeProto.Dim dim1 = TensorShapeProto.Dim.newBuilder().setSize(2).build();
        TensorShapeProto.Dim dim2 = TensorShapeProto.Dim.newBuilder().setSize(9).build();
        TensorShapeProto shape = TensorShapeProto.newBuilder().addDim(dim1).addDim(dim2).build();
        featuresTensorBuilder.setTensorShape(shape);
        TensorProto featuresTensorProto = featuresTensorBuilder.build();

        // Generate gRPC request
        com.google.protobuf.Int64Value version = com.google.protobuf.Int64Value.newBuilder().setValue(modelVersion).build();
        Model.ModelSpec modelSpec = Model.ModelSpec.newBuilder().setName(modelName).setVersion(version).build();
        Predict.PredictRequest request = Predict.PredictRequest.newBuilder().setModelSpec(modelSpec).putInputs("features", featuresTensorProto).putInputs("keys", keysTensorProto).build();

        // Request gRPC server
        Predict.PredictResponse response;
        try {
            response = blockingStub.predict(request);
            java.util.Map<java.lang.String, org.tensorflow.framework.TensorProto> outputs = response.getOutputs();
            for (java.util.Map.Entry<java.lang.String, org.tensorflow.framework.TensorProto> entry : outputs.entrySet()) {
                System.out.println("Response with the key: " + entry.getKey() + ", value: " + entry.getValue());
            }
        } catch (StatusRuntimeException e) {
            logger.log(Level.WARNING, "RPC failed: {0}", e.getStatus());
            return;
        }
    }

}

@@ -2534,88 +2534,8 @@
 r();
-%0A        keysTensorBuilder.setDtype(org.tensorflow.framework.DataType.DT_INT32);
 %0A%0A  
@@ -2738,32 +2738,437 @@
     %7D%0A%0A        %7D
+%0A%0A        TensorShapeProto.Dim keysDim1 = TensorShapeProto.Dim.newBuilder().setSize(2).build();%0A        TensorShapeProto.Dim keysDim2 = TensorShapeProto.Dim.newBuilder().setSize(1).build();%0A        TensorShapeProto keysShape = TensorShapeProto.newBuilder().addDim(keysDim1).addDim(keysDim2).build();%0A        keysTensorBuilder.setDtype(org.tensorflow.framework.DataType.DT_INT32).setTensorShape(keysShape);
 %0A        TensorP
@@ -3520,92 +3520,8 @@
 r();
-%0A        featuresTensorBuilder.setDtype(org.tensorflow.framework.DataType.DT_FLOAT);
 %0A%0A  
@@ -3780,17 +3780,25 @@
 oto.Dim 
-d
+featuresD
 im1 = Te
@@ -3878,17 +3878,25 @@
 oto.Dim 
-d
+featuresD
 im2 = Te
@@ -3972,17 +3972,25 @@
 peProto 
-s
+featuresS
 hape = T
@@ -4025,17 +4025,25 @@
 .addDim(
-d
+featuresD
 im1).add
@@ -4046,17 +4046,25 @@
 .addDim(
-d
+featuresD
 im2).bui
@@ -4099,16 +4099,69 @@
 Builder.
+setDtype(org.tensorflow.framework.DataType.DT_FLOAT).
 setTenso
@@ -4167,17 +4167,25 @@
 orShape(
-s
+featuresS
 hape);%0A 

Add dim in gRPC client request for java_predict_client