//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2016 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
#include "Cuda.hpp"
#include "ImageCuda.hpp"

#include <set>
#include <vector>

using std::cout;
using std::cerr;
using std::endl;
using std::string;

namespace mgx 
{
  // ---------- Initialization routines ----------
  
#ifdef THRUST_BACKEND_CUDA
  // Check available memory
  size_t memAvail()
  {
    size_t free = 0, total = 0;
    cudaMemGetInfo(&free, &total);
    return (free);
  }
  
  static int cudaDevice = -1;
  static size_t holdMemSize = 0;
  static void** holdmem = 0;
  static int cudaCores = 128;
  static int cudaMaxThreads = 1024;
  
  // Free hold memory
  void freeMem()
  {
    cudaThreadSynchronize();
    cudaError_t cuerr = cudaGetLastError();
    if(cuerr != cudaSuccess)
      cerr << "Cuda freeMem, error before freeing memory:" << cuerr << ":" << cudaGetErrorString(cuerr) << endl;
    if(holdmem != 0)
      cudaFree(holdmem);
    holdmem = 0;
    cudaThreadSynchronize();
    // cudaError_t
    cuerr = cudaGetLastError();
    if(cuerr != cudaSuccess)
      cerr << "Cuda freeMem, error freeing memory:" << cuerr << ":" << cudaGetErrorString(cuerr) << endl;
  }
  
  // Allocate hold memory
  void holdMem()
  {
    cudaError_t cuerr;
    cudaThreadSynchronize();
    if(holdmem != NULL) {
      cudaFree(holdmem);
      cuerr = cudaGetLastError();
      if(cuerr != cudaSuccess)
        cerr << "Cuda holdMem, error freeing memory:" << cuerr << ":" << cudaGetErrorString(cuerr) << endl;
      holdmem = NULL;
    }
    static void** new_mem = NULL;
    // The step MUST be smaller or equal to the stop condition or there will be infinite loops. size_t is unsigned, meaning that subtraction is dangerous
    for (size_t trymem = holdMemSize; trymem >= 8 * MByte; trymem -= 4 * MByte) {
      cudaMalloc((void**)&new_mem, trymem);
      cuerr = cudaGetLastError();
      if (cuerr == cudaSuccess) {
        if(holdMemSize != trymem)
          cerr << "Cuda holdMem, reduced memory to:" << trymem / MByte << "MiB" << endl;
        holdMemSize = trymem;
        holdmem = new_mem;
        return;
      }
    }
    cerr << "Cuda holdMem, cannot allocate 8 MiB, giving up:" << cuerr << ":" << cudaGetErrorString(cuerr) << endl;
  }

  int setHoldMem(uint max)
  {
    size_t oldmax = holdMemSize;
    if(max == 0) {
      if(oldmax == 0)
        holdMemSize = memAvail() / 4;
    } else
      holdMemSize = max * MByte;
    holdMem();
    if(holdMemSize != oldmax)
      if(oldmax != 0)
        cout << "Memory reserved for Cuda changed from: " << oldmax / MByte << " MB to: " << holdMemSize / MByte
             << " MB." << endl;
      else
        cout << "Memory reserved for Cuda: " << holdMemSize / MByte << " MB." << endl;
    return (holdMemSize / MByte);
  }
  
  // Initialize cuda
  int initCuda()
  {
    int count = 0;
    holdmem = 0;
  
    int driverVersion, runtimeVersion;
  
    cudaDriverGetVersion(&driverVersion);
    cudaRuntimeGetVersion(&runtimeVersion);
  
    cout << "Cuda driver version: " << (driverVersion / 1000) << "." << (driverVersion % 1000) << endl
         << "Cuda runtime version: " << (runtimeVersion / 1000) << "." << (runtimeVersion % 1000) << endl;
  
    cudaGetDeviceCount(&count);
    if(count == 0) {
      cerr << "There is no device." << endl;
      return (1);
    }
  
    cudaDeviceProp prop;
    for(cudaDevice = 0; cudaDevice < count; cudaDevice++) {
      if(cudaGetDeviceProperties(&prop, cudaDevice) == cudaSuccess) {
        if(prop.major >= 1) {
          break;
        }
      }
    }
    if(cudaDevice == count) {
      cudaDevice = -1;
      cerr << "There is no device supporting CUDA." << endl;
      return (2);
    }
    int computeCapability = prop.major * 10 + prop.minor;
    int resThreadsPerMulti = 768;
    switch(computeCapability) {
    case 10:
    case 11:
      cudaCores = prop.multiProcessorCount * 8;
      resThreadsPerMulti = 768;
      break;
    case 12:
    case 13:
      cudaCores = prop.multiProcessorCount * 8;
      resThreadsPerMulti = 1024;
      break;
    case 20:
      cudaCores = prop.multiProcessorCount * 32;
      resThreadsPerMulti = 1536;
      break;
    case 21:
      cudaCores = prop.multiProcessorCount * 48;
      resThreadsPerMulti = 1536;
      break;
    case 30:
      cudaCores = prop.multiProcessorCount * 192;
      resThreadsPerMulti = 2048;
      break;
    default:   // 1x cards
      cudaCores = prop.multiProcessorCount * 8;
      resThreadsPerMulti = 768;
      break;
    }
    // Take a guess on the future
    if(computeCapability > 30) {
      cudaCores = prop.multiProcessorCount * 192;
      resThreadsPerMulti = 2048;
    }
    cudaMaxThreads = prop.multiProcessorCount * resThreadsPerMulti * 10;
  
    cout << endl;
    cout << "Cuda capable device found, device " << cudaDevice << ": " << prop.name << endl;
    cout << "                 Compute capability: " << prop.major << "." << prop.minor << endl;
    cout << "                       Total memory: " << prop.totalGlobalMem / (1024 * 1024) << " Mb" << endl;
    cout << "                    MultiProcessors: " << prop.multiProcessorCount << endl;
    cout << "     Res.threads per MultiProcessor: " << resThreadsPerMulti << endl;
    cout << "               Max resident threads: " << prop.multiProcessorCount* resThreadsPerMulti << endl;
    cout << "                         Cuda cores: " << cudaCores << endl;
    cout << "                         Clock rate: " << float(prop.clockRate) / 1000000.0 << " GHz" << endl << endl;
  
    cudaSetDevice(cudaDevice);
  
    /*cudaInit(0);*/
    /*if(cudaDevice > 0)*/
    /*cudaCtxDestroy(ctx);*/
    /*cudaCtxCreate(&ctx, CU_CTX_SCHED_YIELD, cudaDevice);*/
  
    size_t free = 0, total = 0;
    cudaMemGetInfo(&free, &total);
    free /= MByte;
    total /= MByte;
  
    cout << "CUDA initialized, " << free << "MB memory available from " << total << " total" << endl << endl;
  
    setHoldMem(0);
    return (0);
  }

  // Check for Cuda errors
  int checkCudaError(const string& msg)
  {
    cudaThreadSynchronize();
    cudaError_t cuerr = cudaGetLastError();
    if(cuerr != cudaSuccess)
      throw(string(msg + " Cuda error string:" + string(cudaGetErrorString(cuerr))));
    return 0;
  }
  
  // Calculate optimal threads, start small
  uint getThreadCount(uint threadPos, uint totThreads)
  {
    static clock_t starttime;
    static uint numThreads;
    static float maxTime = 0;
  
    // Start with same number of threads as cores
    if(threadPos == 0) {
      numThreads = cudaCores;
      maxTime = 0;
      starttime = clock();
    } else {
      clock_t currtime = clock();
      float secs = float(currtime - starttime) / CLOCKS_PER_SEC;
  #ifdef WIN32
      if(secs > .2)
        Sleep(100);
      else
        Sleep(1);
  #else
      if(secs > .2)
        usleep(100000);
      else
        usleep(100);
  #endif
      starttime = currtime;
      if(secs < .02)
        numThreads *= 2;
      else if(secs > .05)
        numThreads /= 2;
      if(numThreads < (uint)cudaCores)
        numThreads = cudaCores;
      if(numThreads > (uint)cudaMaxThreads)
        numThreads = cudaMaxThreads;
  
      if(secs > maxTime) {
        maxTime = secs;
        //      cout << "getThreadCount:Threads:" << numThreads << " Prev MaxTime:" << maxTime << endl;
      }
    }
    return (min(numThreads, totThreads - threadPos));
  }
#else
  #include <omp.h> 

  // If we do not use Cuda than just define stubs
  size_t memAvail()
  {
    return 1024 * MByte;
  }

  void freeMem() {}
  void holdMem() {} 
  int setHoldMem(uint max) { return max; }

  int initCuda() 
  { 
    cout << endl << "Using OMP host emulation instead of Cuda, cores:" 
         << omp_get_num_procs() << endl << endl;
    return 0;
  }

  extern "C" 
  {  
    int checkCudaError(const string& msg)
    {
      return 0;
    }
  }
  
  // Use number of host cores, omp should always be available
  uint getThreadCount(uint threadPos, uint totThreads)
  {
    return totThreads - threadPos;
  }
#endif

  // Memory for user processes
  size_t userMem()
  {
    size_t avail = memAvail();
  
    // Allow 20% buffer
    return (avail - avail / 5);
    // return(holdMemSize - holdMemSize/10);
  }
  
  // Reduce available user mem
  size_t memLeft(size_t memrq, size_t roomrq, size_t mem)
  {
    if(memrq + roomrq > mem)
      throw(string("memLeft::Not enough cuda memory to allocate min/max list"));
    return (mem - memrq);
  }
  
  // Error message
  int errMsg(const string& s)
  {
    cerr << s << endl;
    return 1;
  }

  // ---------- Externally exported functions ----------
  // Initialize cuda
  cuda_EXPORT
  void initGPU() { initCuda(); }

  cuda_EXPORT
  int setHoldMemGPU(uint max) { return setHoldMem(max); }
} 
