// matrix transpose, from the Thrust package examples

#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/functional.h>
#include <thrust/gather.h>
#include <thrust/scan.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>
#include <iostream>
#include <iomanip>

// convert a linear index to a linear index in the transpose 
struct transpose_index : public thrust::unary_function<size_t,size_t>
{
  size_t m, n;

  __host__ __device__
  transpose_index(size_t _m, size_t _n) : m(_m), n(_n) {}

  __host__ __device__
  size_t operator()(size_t linear_index)
  {
      size_t i = linear_index / n;
      size_t j = linear_index % n;

      return m * j + i;
  }
};

// convert a linear index to a row index
struct row_index : public thrust::unary_function<size_t,size_t>
{
  size_t n;

  __host__ __device__
  row_index(size_t _n) : n(_n) {}

  __host__ __device__
  size_t operator()(size_t i)
  {
      return i / n;
  }
};

// transpose an M-by-N array
template <typename T>
void transpose(size_t m, size_t n, thrust::device_vector<T>& src, thrust::device_vector<T>& dst)
{
  thrust::counting_iterator<size_t> indices(0);

  thrust::gather
    (thrust::make_transform_iterator(indices, transpose_index(n, m)),
     thrust::make_transform_iterator(indices, transpose_index(n, m)) + dst.size(),
     src.begin(),
     dst.begin());
}

void transp(int *inmat, int *outmat, int nr, int nc)
{
   thrust::device_vector<int> dmat(inmat,inmat+nr*nc);
   // make space for the transpose
   thrust::device_vector<int> dxp(nr*nc);
   transpose(nr,nc,dmat,dxp);
   thrust::copy(dxp.begin(),dxp.end(),outmat);
}

int rand16()  // generate random integers mod 16
{  return rand() % 16;  }

int main(int argc, char **argv)
{  int nr = atoi(argv[1]); int nc = nr;
   int *mat = (int *) malloc(nr*nc*sizeof(int));
   int *matxp = (int *) malloc(nr*nc*sizeof(int));
   thrust::generate(mat,mat+nr*nc,rand16);
//    thrust::copy(mat, mat+nr*nc,
//       std::ostream_iterator<int>(std::cout, " ")); 
//    std::cout << "\n"; 
   transp(mat,matxp,nr,nc);
//    thrust::copy(matxp, matxp+nr*nc,
//       std::ostream_iterator<int>(std::cout, " ")); 
//    std::cout << "\n"; 
}