#include <paddle/capi.h>
#include <time.h>

#include "../common/common.h"

// Modify this path as needed.
#define CONFIG_BIN "./trainer_config.bin"
// Modify this path as needed.
// This demo assumes that merged model is not used, then this path is the
// directory storing all the trained parameters.
// If the model is trained by PaddlePaddle V2 API, the model is saved as
// a compressed file. You need to uncompress the compressed file first.
#define MODEL_PATH "models/pass_4"

int main() {
  // Initalize the PaddlePaddle runtime environment.
  char* argv[] = {"--use_gpu=False"};
  CHECK(paddle_init(1, (char**)argv));

  // Read the binary configuration file generated by `convert_protobin.sh`
  long size;
  void* buf = read_config(CONFIG_BIN, &size);

  // Create the gradient machine for inference.
  paddle_gradient_machine machine;
  CHECK(paddle_gradient_machine_create_for_inference(&machine, buf, (int)size));

  // Load the trained model. Modify the parameter MODEL_PATH to set the correct
  // path of the trained model.
  CHECK(paddle_gradient_machine_load_parameter_from_disk(machine, MODEL_PATH));

  // Inputs and outputs of the network are organized as paddle_arguments object
  // in C-API. In the comments below, "argument" specifically means one input of
  // the neural network in PaddlePaddle C-API.
  paddle_arguments in_args = paddle_arguments_create_none();

  // There is only one data layer in this demo MNIST network, invoke this
  // function to create one argument.
  CHECK(paddle_arguments_resize(in_args, 1));

  // Each argument needs one matrix or one ivector (integer vector, for sparse
  // index input, usually used in NLP task) to holds the real input data.
  // In the comments below, "matrix" specifically means the object needed by
  // argument to hold the data. Here we create the matrix for the above created
  // agument to store the testing samples.
  paddle_matrix mat =
      paddle_matrix_create(/* height = batch size */ 1,
                           /* width = dimensionality of the data layer */ 784,
                           /* whether to use GPU */ false);

  paddle_real* array;
  // Get the pointer pointing to the start address of the first row of the
  // created matrix.
  CHECK(paddle_matrix_get_row(mat, 0, &array));

  // Fill the matrix with a randomly generated test sample.
  srand(time(0));
  for (int i = 0; i < 784; ++i) {
    array[i] = rand() / ((float)RAND_MAX);
  }

  // Assign the matrix to the argument.
  CHECK(paddle_arguments_set_value(in_args, 0, mat));

  // Create the output argument.
  paddle_arguments out_args = paddle_arguments_create_none();

  // Invoke the forward computation.
  CHECK(paddle_gradient_machine_forward(machine,
                                        in_args,
                                        out_args,
                                        /* is train taks or not */ false));

  // Create the matrix to hold the forward result of the neural network.
  paddle_matrix prob = paddle_matrix_create_none();
  // Access the matrix of the output argument, the predicted result is stored in
  // which.
  CHECK(paddle_arguments_get_value(out_args, 0, prob));

  uint64_t height;
  uint64_t width;
  CHECK(paddle_matrix_get_shape(prob, &height, &width));
  CHECK(paddle_matrix_get_row(prob, 0, &array));

  printf("Prob: \n");
  for (int i = 0; i < height * width; ++i) {
    printf("%.4f ", array[i]);
    if ((i + 1) % width == 0) {
      printf("\n");
    }
  }
  printf("\n");

  // The cleaning up.
  CHECK(paddle_matrix_destroy(prob));
  CHECK(paddle_arguments_destroy(out_args));
  CHECK(paddle_matrix_destroy(mat));
  CHECK(paddle_arguments_destroy(in_args));
  CHECK(paddle_gradient_machine_destroy(machine));

  return 0;
}