// matrix transpose, for_each version

#include <stdio.h>
#include <thrust/device_vector.h>

// functor; holds iterators for the input and output matrices, and each
// invocation of the function copies from one element from the former to
// the latter
struct copyelt2xp  
{
   int nrow;  
   int ncol;  
   const  thrust::device_vector<int>::iterator m;  // input matrix 
   const  thrust::device_vector<int>::iterator mxp;  // output matrix
   int *m1,*mxp1;
   copyelt2xp(thrust::device_vector<int>::iterator _m, 
            thrust::device_vector<int>::iterator _mxp, 
            int _nr, int _nc): 
      m(_m), mxp(_mxp), nrow(_nr), ncol(_nc) {
         m1 = thrust::raw_pointer_cast(&m[0]);
         mxp1 = thrust::raw_pointer_cast(&mxp[0]);
      }
   __device__
   void operator()(const int i)  
   // copies the i-th element of the input matrix to the output matrix
   {  // elt i in input is row r, col c there
      int r = i / ncol; int c = i % ncol;  
      // that elt will be row c and col r in output, which has nrow
      // cols, so copy as follows
      mxp1[c*nrow+r] = m1[r*ncol+c];
   }
};

// transpose nr x nc inmat to outmat
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);
   thrust::counting_iterator<int> seqb(0);
   thrust::counting_iterator<int> seqe = seqb + nr*nc;
   // for each i in seq, copy the matrix elt to its spot in the
   // transpose
   // thrust::for_each(seqb,seqe,
   //    copyelt2xp(dmat.begin(),dxp.begin(),nr,nc));
   copyelt2xp z(dmat.begin(),dxp.begin(),nr,nc);
   thrust::for_each(seqb,seqe, z);
   thrust::copy(dxp.begin(),dxp.end(),outmat);
}

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

// test code:  cmd line args are matrix size, then row, col of elt to be
// checked

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);
   int checkrow = atoi(argv[2]);
   int checkcol = atoi(argv[3]);
   printf("%d\n",mat[checkrow*nc+checkcol]);
   transp(mat,matxp,nr,nc);
   printf("%d\n",matxp[checkcol*nc+checkrow]);
}