#include <cstdio>
#include <cstdlib>
#include <unistd.h>

#include "uNet.h"
#include "util.h"
#include <cuda_runtime.h>

const int WARM_UP = 2;
const int MEASURE = 3;

// 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, 640, 959});
  Tensor *output = new Tensor({N, 2, 640, 959});
  
  size_t input_size = 0;
  read_binary((void*)input->buf, input_fname, &input_size);

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

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

  // 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;
	printf("Done! (%lf img/sec)\n", N/elapsed_time/MEASURE);

  write_binary((void*)output->buf, output_fname, (size_t)(N * 2 * 640 * 959));

  printf(" Writing to %s ...", output_fname);
  fflush(stdout);
  
  printf("Done!\n\n");

  // Finalize program
  uNet_finalize();
}