/* ========================================================================== *
 *  Utilities.
 *
 *  Author: Firas Safadi (fsafadi@ulg.ac.be)
 * ========================================================================== */

#pragma once

#include <iostream>
#include <fstream>
#include <sstream>
#include <iomanip>
#include <algorithm>
#include <limits>

#include <intrin.h>
#include <conio.h>
#include <windows.h>

#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "curand_kernel.h"
#include "cublas_v2.h"
#include "helper_cuda.h"

#pragma intrinsic(__rdtsc)

//Uncomment the following line to switch to 64-bit precision
//#define DOUBLE_PRECISION

#define LOG_FLT(f) ((((*((__int32*)&(f)) >> 23) & 0xFF) - 0x7F) / 3.3219281)
#define LOG_DBL(d) ((((*((__int64*)&(d)) >> 52) & 0x7FF) - 0x3FF) / 3.3219281)

#define ODD(n) ((n) & 1)

#define MIN(a, b) ((a) < (b) ? (a) : (b))
#define RDIV(a, b) ((a) % (b) ? (a) / (b) + 1 : (a) / (b))

#define RDTSC() __rdtsc()

#define BUF_SIZE 16

#define AssertKernelLaunch() cuda_call(cudaPeekAtLastError())

#define AsyncKernelLaunch(kernel_name, grid_size, block_size, ...) \
  do { kernel_name <<< (grid_size), (block_size) >>> (__VA_ARGS__); cuda_call(cudaPeekAtLastError()); } while (0)
#define SyncKernelLaunch(kernel_name, grid_size, block_size, ...) \
  do { kernel_name <<< (grid_size), (block_size) >>> (__VA_ARGS__); cuda_call(cudaPeekAtLastError()); cudaDeviceSynchronize(); } while (0)

typedef unsigned int nat;
typedef size_t bytes;

#ifndef DOUBLE_PRECISION
typedef float real;
#else
typedef double real;
#endif

typedef enum { DARKBLUE = 1, DARKGREEN, DARKTEAL, DARKRED, DARKPINK, DARKYELLOW, GRAY, DARKGRAY, BLUE, GREEN, TEAL, RED, PINK, YELLOW, WHITE } Color;

static const char* binary_units[] = { "bytes", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB" };

static HANDLE console_handle = NULL;

inline void set_color(Color c)
{
  if (console_handle == NULL)
    console_handle = GetStdHandle(STD_OUTPUT_HANDLE);

  SetConsoleTextAttribute(console_handle, c);
}

inline void flush_input()
{
  while (_kbhit())
    _getch();
}

inline bool yes_no(const char* msg)
{
  flush_input();

  char q;

  do {
    std::cout << std::endl;
    std::cout << msg << " (y/n) ";
    q = _getch();
  } while (q != 'y' && q != 'n');

  std::cout << std::endl;

  return q == 'y';
}

inline void pause()
{
  flush_input();

  std::cout << std::endl;
  std::cout << "Press any key to continue...";

  _getch();

  std::cout << std::endl;
}

inline void pause(const char* msg)
{
  std::cout << std::endl;
  std::cout << msg << std::endl;

  pause();
}

inline void error(const char* msg)
{
  pause(msg);

  exit(EXIT_FAILURE);
}

inline void cuda_call(cudaError_t x)
{
  if (x != cudaSuccess)
    error(_cudaGetErrorEnum(x));
}

inline void cublas_call(cublasStatus_t x)
{
  if (x != CUBLAS_STATUS_SUCCESS)
    error(_cudaGetErrorEnum(x));
}

inline void curand_call(curandStatus_t x)
{
  if (x != CURAND_STATUS_SUCCESS)
    error(_cudaGetErrorEnum(x));
}

inline int cuda_cores(const cudaDeviceProp& deviceProp)
{
  return _ConvertSMVer2Cores(deviceProp.major, deviceProp.minor) * deviceProp.multiProcessorCount;
}

inline char* size_to_string(char* dest, bytes n)
{
  bytes a = n;

  nat i;

  for (i = 0; a > 1024; i++)
    a >>= 10;

  std::stringstream s;

  s << (i > 0 ? n >> 10 * i : n) << " " << binary_units[i];

  strncpy(dest, s.str().c_str(), BUF_SIZE);

  return dest;
}

inline bool approx_equal(float a, float b, nat bits)
{
  __int32* p = (__int32*)&a;
  __int32* q = (__int32*)&b;

  return *p >> (23 - bits) == *q >> (23 - bits);
}

inline bool approx_equal(double a, double b, nat bits)
{
  __int64* p = (__int64*)&a;
  __int64* q = (__int64*)&b;

  return *p >> (52 - bits) == *q >> (52 - bits);
}

inline real rand(real a, real b)
{
  return (b - a)*((real)rand() / (real)RAND_MAX) + a;
}

class Interval
{
public:
  Interval(real a = 0.0, real b = 1.0) : low(_low), high(_high) { _low = a; _high = b; }
  ~Interval() {}

  void set(real a, real b) { if (a > b) error("Lower bound cannot be greater than upper bound"); _low = a; _high = b; }

  const real& low;
  const real& high;

private:
  real _low;
  real _high;
};

class Timer
{
public:
  Timer() { _start = 0; _stop = 0; }
  ~Timer() {}

  void start() { timeBeginPeriod(1); _start = timeGetTime(); }
  void stop() { _stop = timeGetTime(); timeEndPeriod(1); }
  
  DWORD msec() { return _stop - _start; }
  float sec() { return msec() / 1000.0; }

private:
  nat _start;
  nat _stop;
};
