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

#define CONFIG_BIN "./trainer_config.bin"

int main() {
  // Initalize Paddle
  char* argv[] = {"--use_gpu=False"};
  CHECK(paddle_init(1, (char**)argv));

  // Reading config binary file. It is generated by `convert_protobin.sh`
  long size;
  void* buf = read_config(CONFIG_BIN, &size);

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

  // Loading parameter. Uncomment the following line and change the directory.
  // CHECK(paddle_gradient_machine_load_parameter_from_disk(machine,
  //                                                "./some_where_to_params"));
  paddle_arguments in_args = paddle_arguments_create_none();

  // There is only one input of this network.
  CHECK(paddle_arguments_resize(in_args, 1));

  // Create input ids.
  int sentence_ids[] = {83, 48, 20, 84, 394, 853, 64, 53, 64};

  paddle_ivector sentence = paddle_ivector_create(
      sentence_ids, sizeof(sentence_ids) / sizeof(int), false, false);
  CHECK(paddle_arguments_set_ids(in_args, 0, sentence));

  int seq_pos_array[] = {0, sizeof(sentence_ids) / sizeof(int)};

  paddle_ivector seq_pos = paddle_ivector_create(
      seq_pos_array, sizeof(seq_pos_array) / sizeof(int), false, false);

  CHECK(paddle_arguments_set_sequence_start_pos(in_args, 0, 0, seq_pos));

  paddle_arguments out_args = paddle_arguments_create_none();
  CHECK(paddle_gradient_machine_forward(machine,
                                        in_args,
                                        out_args,
                                        /* isTrain */ false));
  paddle_matrix prob = paddle_matrix_create_none();

  CHECK(paddle_arguments_get_value(out_args, 0, prob));

  paddle_real* array;

  CHECK(paddle_matrix_get_row(prob, 0, &array));

  printf("Prob: ");
  for (int i = 0; i < 2; ++i) {
    printf("%.2f ", array[i]);
  }
  printf("\n");

  CHECK(paddle_matrix_destroy(prob));
  CHECK(paddle_arguments_destroy(out_args));
  CHECK(paddle_ivector_destroy(seq_pos));
  CHECK(paddle_ivector_destroy(sentence));
  CHECK(paddle_arguments_destroy(in_args));
  CHECK(paddle_gradient_machine_destroy(machine));

  return 0;
}