// includes, system
#include <iostream>
#include <stdlib.h>
#include <time.h>

// Utilities and system includes
#include <cuda_runtime.h>
#include <curand.h>

#include <helper_functions.h>
#include <helper_cuda.h>

#include <curand.h>
#include <cublas_v2.h>

#include <thrust/device_vector.h>
#include <thrust/device_ptr.h>
#include <thrust/host_vector.h>
#include <thrust/transform.h>

void vector_axpy(const double a,const double *X, const double *Y, double *Z, int N){
  for (int i=0;i<N;i++)
    Z[i]=a*X[i]+Y[i];
}

double errorinf(const double *X,const double *Y,int N){
  double norm=fabs(X[0]-Y[0]);
  for (int i=1;i<N;i++){
    double s=fabs(X[i]-Y[i]);
    if (s>norm) norm=s;
  }
  return norm;
}

template <typename T>
struct trans : public thrust::unary_function<T,T>
{
    T a, b;

    __host__ __device__
    trans(T _a, T _b) : a(_a), b(_b) {}

    __host__ __device__
    T operator()(T x)
    {
        return a+(b-a)*x;
    }
};

thrust::device_vector<double> UniformRandom(
  int N, double a, double b, unsigned int mainSeed
){
  thrust::device_vector<double> V(N);
  curandGenerator_t m_prng;
  //Create a new generator
  curandCreateGenerator(&m_prng, CURAND_RNG_PSEUDO_DEFAULT);
  //Set the generator options
  curandSetPseudoRandomGeneratorSeed(m_prng, (unsigned long) mainSeed);
  curandGenerateUniformDouble(m_prng, thrust::raw_pointer_cast(V.data()), N);
  thrust::transform(V.begin(),V.end(),V.begin(), trans<double>(a,b));
  curandDestroyGenerator(m_prng);
  return V;
}

void print(const thrust::device_vector<double>  &V,int n){
  int N=V.size();
  int nd=N-n,np=n;
  
  if (n>N){
    np=N;nd=N+1;
  }
  std::cout << "  <device_vector>["<< V.size() << "]\n";
  //std::setprecision(16)
  std::cout.precision(16);  
  for (int i=0;i<np;i++)
    std::cout << " [" << i << "]: " << V[i] << std::endl;
  if (nd<N)
    std::cout << " ...\n";
  for (int i=nd;i<N;i++)
    std::cout << " [" << i << "]: " << V[i] << std::endl;
  std::cout << "\n";
}

void print(const thrust::host_vector<double>  &V,int n){
  int N=V.size();
  int nd=N-n,np=n;
  
  if (n>N){
    np=N;nd=N+1;
  }
  std::cout << "  <host_vector>["<< V.size() << "]\n";
  //std::setprecision(16)
  std::cout.precision(16);  
  for (int i=0;i<np;i++)
    std::cout << " [" << i << "]: " << V[i] << std::endl;
  if (nd<N)
    std::cout << " ...\n";
  for (int i=nd;i<N;i++)
    std::cout << " [" << i << "]: " << V[i] << std::endl;
  std::cout << "\n";
}
#define CURAND_CALL(x) do { if((x)!=CURAND_STATUS_SUCCESS) { \
    printf("Error at %s:%d\n",__FILE__,__LINE__);\
    return EXIT_FAILURE;}} while(0)

int main(int argc, char** argv){
  int N=10000;
  thrust::host_vector<double> h_X(N),h_Y(N),h_Z(N),h_Zgpu(N);
  thrust::device_vector<double> d_X,d_Y,d_Z;
  
  cublasHandle_t handle;
  cublasStatus_t stat = cublasCreate(&handle);
  if (stat != CUBLAS_STATUS_SUCCESS) {
    printf ("CUBLAS initialization failed\n");
    return EXIT_FAILURE;
  }
  time_t t1;
  time(&t1);
  srand48((long) t1); // Initialize
  
  std::cout<< "A) Generate uniform random [-2,2] on <device>  : \nd_X =\n";
  d_X=UniformRandom(N,-2.,2.,(long) t1);
  print(d_X,3);
  
  std::cout<< "B) Generate uniform random ]0,1] on <host>     : \nh_X =\n"; 
  thrust::generate(h_X.begin(), h_X.end(), drand48);
  print(h_X,3);
  
  std::cout<< "C) Generate uniform random [0,1] on <host>     : \nh_Y =\n"; 
  thrust::generate(h_Y.begin(), h_Y.end(), drand48);
  print(h_Y,3);
  
  std::cout<< "D) Copy <host> array h_x on <device> array d_X : \nd_X =\n";
  d_X=h_X;
  print(d_X,3);
  std::cout<< "E) Copy <host> array h_y on <device> array d_Y : \nd_Y =\n";
  d_Y=h_Y;
  print(d_Y,3);
  
  std::cout<< "F) Compute on <device>,   d_Y <- d_Y - d_X     : \nd_Y =\n";
  double alpha=-1.0;
  stat=cublasDaxpy(handle,N,&alpha,thrust::raw_pointer_cast(d_X.data()),1,
		                   thrust::raw_pointer_cast(d_Y.data()),1);
  if (stat != CUBLAS_STATUS_SUCCESS) {
    printf ("CUBLAS cublasDaxpy failed\n");
    return EXIT_FAILURE;
  }
  print(d_Y,3);
  
  std::cout<< "G) Copy <device> array d_Y on <host> array h_Zgpu : \nh_Zgpu =\n";
  h_Zgpu=d_Y;
  print(h_Zgpu,3);
  
  std::cout<< "H) Compute h_Z <- h_Y - h_X                  : \nh_Z =\n";
  vector_axpy(-1.,&h_X[0], &h_Y[0], &h_Z[0], N);
  print(h_Z,3);
  
  std::cout<< "I) Compute    max|h_Zgpu - h_Z| = " << errorinf(&h_Zgpu[0],&h_Z[0],N);
  std::cout << "\n...End\n";
}