#include <cuda_runtime.h>
#include <unistd.h>

#include <cstdio>
#include <cstdlib>

#include "uNet.h"
#include "util.h"

const int WARM_UP = 2;

#ifdef TEST
const int MEASURE = 1;
#else
const int MEASURE = 3;
#endif

// Global variables
int N = 1;
char *parameter_fname;
char *output_fname;
char *input_fname;

int main(int argc, char **argv) {
  check_and_parse_args(argc, argv);
  print_model();

  // Initialize model
  uNet_initialize(N, parameter_fname);

  Tensor *input = new Tensor({N, 3, 128, 191});
  Tensor *output = new Tensor({N, 2, 128, 191});

  size_t input_size = 0;
  read_binary((void *) input->buf, input_fname, &input_size);

  printf(" process %d image(s)...", N);
  fflush(stdout);

// warm_up
#ifndef TEST
  printf("\nWarmimg up.");
  for (int i = 0; i < WARM_UP; ++i) {
    uNet(input, output, N);
    printf(".");
  }
  cudaDeviceSynchronize();
  printf("\n");
#endif

  // run uNet and measure time
  printf("\nProcess.");
  double uNet_st = get_time();
  for (int j = 0; j < MEASURE; ++j) {
    uNet(input, output, N);
    printf(".");
  }
  cudaDeviceSynchronize();
  double uNet_en = get_time();
  printf("\n");
  double elapsed_time = uNet_en - uNet_st;
  elapsed_time = elapsed_time / MEASURE;
  printf("Done! %lfsec (%lf img/sec)\n", elapsed_time, N / elapsed_time);

  write_binary((void *) output->buf, output_fname, (size_t)(N * 2 * 128 * 191));

  printf(" Writing to %s ...", output_fname);
  fflush(stdout);

  printf("Done!\n\n");

  // Finalize program
  uNet_finalize();
}