//
// 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.
// 
#ifndef DIST_MATRIX_CUDA_HPP
#define DIST_MATRIX_CUDA_HPP

#include <Cuda.hpp>
#include <ThrustTypes.hpp>
#include <CachedAlloc.hpp>

#include <thrust/inner_product.h> 
#include <thrust/extrema.h> 

namespace mgx
{
  // Vector (or matrix)-scalar multiply
  template<typename T>
  struct multVSOP
  {
    const double a;
  
    multVSOP(double _a) : a(_a) {}
  
    __host__ __device__
    T operator()(const T &x) const 
    { 
      return a * x;
    }
  };
  
  // Matrix-vector multiply
  template<typename MatT, typename VecT>
  struct multMVOP
  {
    uint *Nb;
    MatT *Mv, *Me;
    VecT *v;
    VecT *r;
    uint nbs;

    // Types of deice vectors
    typedef typename thrust::device_vector<MatT> DevMatV;
    typedef typename thrust::device_vector<VecT> DevVecV;

    multMVOP(DeviceVu *_Nb, DevMatV *_Mv, DevMatV *_Me, DevVecV *_v, DevVecV *_r, uint _nbs) 
     : Nb(devP(_Nb)), Mv(devP(_Mv)), Me(devP(_Me)), v(devP(_v)), r(devP(_r)), nbs(_nbs) {}
  
    __host__ __device__
    void operator()(const uint vtx) const 
    { 
      uint edg = vtx * nbs;
      VecT sum = Mv[vtx] * v[vtx];
      
      // Sum over neighbors
      for(uint i = 0; i < nbs; i++, edg++)
        if(Nb[edg] != vtx)
          sum += Me[edg] * v[Nb[edg]];
  
      r[vtx] = sum;
    }
  };
  
  // Add to diagonal 
  template <typename T>
  struct addToDiagOP
  {
    T *Mv;
    double a;
    int sz;

    typedef typename thrust::device_vector<T> DevMatV;
  
    addToDiagOP(DevMatV *_Mv, double _a)
     : Mv(devP(_Mv)), a(_a) {}
  
     __host__ __device__
    void operator()(const uint vtx) const 
    { 
      T &m = Mv[vtx];
      
      for(int i = 0; i < T::numrows; ++i)
        m[i][i] += a;
    }
  };
  
  // Jacobi pre-conditioner
  template<typename T>
  struct jacobiPreCondOP
  {
    uint *Nb;
    T *Mv, *Me;
    T *Av, *Ae;
    uint nbs;
  
    jacobiPreCondOP(DeviceVu *_Nb, thrust::device_vector<T> *_Mv, thrust::device_vector<T> *_Me, 
       thrust::device_vector<T> *_Av, thrust::device_vector<T> *_Ae, uint _nbs) 
     : Nb(devP(_Nb)), Mv(devP(_Mv)), Me(devP(_Me)), Av(devP(_Av)), Ae(devP(_Ae)), nbs(_nbs) {}
  
     __host__ __device__
    void operator()(const uint vtx) const 
    { 
      T &m = Mv[vtx];
      const T &a = Av[vtx];
      m =  inverse(a);
    }
  };
  
  // Saxpy
  template<typename T>
  struct saxpyOP 
  { 
    const double a; 
    saxpyOP(double _a) : a(_a) {}
  
     __host__ __device__ 
     T operator()(const T &x, const T &y) const 
     { 
       return a * x + y; 
     } 
  }; 

  // Sadly numeric_limits complains about calling host functions
  template<typename T>
  struct limits
  {
    __device__ __host__
    T static min();
    __device__ __host__
    T static max();
  };
  template<> double limits<double>::min() { return DBL_MIN; }
  template<> float limits<float>::min() { return FLT_MIN; }
  template<> double limits<double>::max() { return DBL_MAX; }
  template<> float limits<float>::max() { return FLT_MAX; }

  template<typename T>
  struct minElementOP : public thrust::unary_function<T, typename T::value_type>
  {
    typedef typename T::value_type ValType;

    __host__ __device__
    ValType operator()(T &a)
    {
      ValType min = limits<ValType>::max();
      uint n = sizeof(T)/sizeof(ValType);
      ValType *p = a.data();
      for(uint i = 0; i < n; i++) {
        if(min > *p)
          min = *p;
        ++p;
      }
      
      return min;
    }
  };

  template<typename T>
  struct maxElementOP : public thrust::unary_function<T, typename T::value_type>
  {
    typedef typename T::value_type ValType;

    __host__ __device__
    ValType operator()(T &a)
    {
      ValType max = limits<ValType>::min();
      uint n = sizeof(T)/sizeof(ValType);
      ValType *p = a.data();
      for(uint i = 0; i < n; i++) {
        if(max < *p)
          max = *p;
        ++p;
      }
      
      return max;
    }
  };

  // Component multiply
  template<typename T>
  struct multCompOP: public thrust::binary_function<T, T, T>
  {
     __host__ __device__

    T operator()(const T &x,const T &y)
    {
	    T v;
      for(uint i = 0; i < T::numElems; ++i)
	      v[i] = x[i] * y[i];

      return v;
    }
  };

  template<typename T>
  struct addOP : public thrust::binary_function<T, T, T>
  {
    __host__ __device__

    T operator()(const T &x,const T &y)
    {
	    T v = x + y;

      return v;
    } 
  };

  //
  // Distributed matrix GPU calls
  //
  template <typename T>
  int multGPU(T *v, typename T::value_type::value_type a, T *r)
  {
    if(!r or !v)
      return 1;

    thrust::transform(v->begin(), v->end(), r->begin(), multVSOP<typename T::value_type>(a));

    return 0;
  }

  // Matrix or vector inner product
  template<typename T>
  int multGPU(T *v1, T *v2, typename T::value_type::value_type &r)
  {
    if(!v1 or !v2)
      return 1;

    typename T::value_type init;
    for(uint i = 0; i < T::value_type::numElems; ++i)
      init[i] = 0;

// Breaks in Cuda >7
#ifdef THRUST_BACKEND_CUDA 
    typename T::value_type res = thrust::inner_product(thrust::cuda::par(mgx::cachedAlloc), v1->begin(), v1->end(), 
                        v2->begin(), init, addOP<typename T::value_type>(), multCompOP<typename T::value_type>());
#else
    typename T::value_type res = thrust::inner_product(v1->begin(), v1->end(), 
                        v2->begin(), init, addOP<typename T::value_type>(), multCompOP<typename T::value_type>());
#endif
    r = 0;
    for(uint i = 0; i < T::value_type::numElems; ++i)
      r += res[i];

    return(0);
  }

  // Matrix-vector multiply
  template <typename TM, typename TV>
  int multGPU(DeviceVu *nb, TM *mv, TM *me, TV *v, TV *r)
  {
    if(!nb or !r or !mv or !me or !v)
      return 1;

    // Grab correctly typed device pointers
    uint nbs = me->size()/mv->size();
    uint n = nb->size()/nbs;
  
    thrust::counting_iterator<int, thrust::device_system_tag> first(0);
    thrust::counting_iterator<int, thrust::device_system_tag> last(n);
    thrust::for_each(first, last, multMVOP<typename TM::value_type, typename TV::value_type>(nb, mv, me, v, r, nbs));
  
    return(0);
  }

  // Add to diagonal
  template<typename T>
  int addToDiagGPU(typename T::value_type::value_type a, T *r)
  {
    if(!r)
      return 1;

    thrust::counting_iterator<int, thrust::device_system_tag> first(0);
    thrust::counting_iterator<int, thrust::device_system_tag> last(r->size());
    thrust::for_each(first, last, addToDiagOP<typename T::value_type>(r, a));
  
    return(0);
  }

  // Matrix or vector addition
  template<typename T>
  int addGPU(T *v1, T *v2, T *r)
  {
    if(!r or !v1 or !v2)
      return 1;

    thrust::transform(v1->begin(), v1->end(), v2->begin(), r->begin(), thrust::plus<typename T::value_type>());
  
    return(0);
  } 

  // Matrix or vector subtraction
  template<typename T>
  int subtractGPU(T *v1, T *v2, T *r)
  {
    if(!r or !v1 or !v2)
      return 1;

    thrust::transform(v1->begin(), v1->end(), v2->begin(), r->begin(), thrust::minus<typename T::value_type>());
  
    return(0);
  }  

  // Perform saxpy r = a * v1 + v2
  template<typename T>
  int saxpyGPU(T *v1,  T *v2, typename T::value_type::value_type a, T *r)
  {
    if(!v1 or !v2)
      return 1;

    thrust::transform(v1->begin(), v1->end(), v2->begin(), r->begin(), saxpyOP<typename T::value_type>(a));

    return(0);
  } 

  // Fill matrix or vector with scalar
  template<typename T>
  int fillGPU(typename T::value_type::value_type a, T *r)
  {
    if(!r)
      return 1;

    typename T::value_type fill(a);

    thrust::fill(r->begin(), r->end(), fill);
  
    return(0);
  } 

  // Find min value
  template<typename T>
  int minGPU(T *v, typename T::value_type::value_type &r)
  {
    if(!v)
      return 1;

    thrust::device_vector<typename T::value_type::value_type> s(v->size());
    thrust::transform(v->begin(), v->end(), s.begin(), minElementOP<typename T::value_type>());

#ifdef THRUST_BACKEND_CUDA 
    r = *thrust::min_element(thrust::cuda::par(mgx::cachedAlloc), s.begin(), s.end());
#else
    r = *thrust::min_element(s.begin(), s.end());
#endif 

    return(0);
  }

  // Find max value
  template<typename T>
  int maxGPU(T *v, typename T::value_type::value_type &r)
  {
    if(!v)
      return 1;

    thrust::device_vector<typename T::value_type::value_type> s(v->size());
    thrust::transform(v->begin(), v->end(), s.begin(), maxElementOP<typename T::value_type>());

#ifdef THRUST_BACKEND_CUDA 
    r = *thrust::max_element(thrust::cuda::par(mgx::cachedAlloc), s.begin(), s.end());
#else
    r = *thrust::max_element(s.begin(), s.end());
#endif 

    return(0);
  }

  // Jacobi preconditioner
  template<typename T>
  int jacobiPreCondGPU(DeviceVu *nb, T *mv, T *me, T *av, T *ae)
  {
    if(!nb or !mv or !me or !av or !ae)
      return 1;

    // Grab correctly typed device pointers
    uint nbs = me->size()/mv->size();
    uint n = nb->size()/nbs;

    thrust::counting_iterator<int, thrust::device_system_tag> first(0);
    thrust::counting_iterator<int, thrust::device_system_tag> last(n);
    thrust::for_each(first, last, jacobiPreCondOP<typename T::value_type>(nb, mv, me, av, ae, nbs));
  
    return(0);
  }
}
#endif
