using namespace std;

#include <iostream>
#include <fstream>
#include <cassert>
#include <cstdlib>
#include <openacc.h>
#include <cuda_runtime.h>

    typedef char*   pickle_t; 
    typedef char    pickle_v; 

    template < typename type > void subv( int n1, int n2, type *a, type *as[] )
   {
       int i;
       #pragma acc kernels present(a[0:n1*n2],as[0:n1])
      {
          as[0]= a;
         #pragma acc loop seq
          for( i=1;i<n1;i++ )
         {
             as[i]= as[i-1]+n2;
         }
      }
   }
 
    template < typename type > void subv_cpu( int n1, int n2, type *a, type *as[] )
   {
       int i;
      {
          as[0]= a;
          for( i=1;i<n1;i++ )
         {
             as[i]= as[i-1]+n2;
         }
      }
   }

    template <typename type> void pckle( size_t *len, const int n, int *iref, const type *var, pickle_t buf )
   {
       type*                                  tmp=NULL;
       int                                    i,j;
       size_t                                   l= (size_t)(n*sizeof(*var));
       tmp= new type[n];
      #pragma acc enter data create(tmp[0:n])
      #pragma acc parallel loop \
       present(var[0:n],tmp[0:n],iref[0:n]) \
       default(none)
       for( j=0;j<n;j++ )
      {
          i= iref[j];
          tmp[j]= var[i];
      }
       #pragma acc host_data use_device(buf,tmp) 
      {
          cudaMemcpy( (void*)(buf+(*len)),  (const void*)(tmp),        l, cudaMemcpyDeviceToDevice  );
      }
       (*len)+= l;
      //delete[] tmp;
       #pragma acc exit data copyout(tmp[0:n])
   }

    template <typename type> void unpckle( size_t *len, int n, type *var, const pickle_t buf )
   {
       size_t l= (size_t)(n*sizeof(*var));
      #pragma acc host_data use_device(buf) 
      {
          cudaMemcpy( (void*)var, (const void*)(buf+(*len)), l, cudaMemcpyDeviceToDevice ); 
      }
       (*len)+= l;
   }

    int main(int argc, char** argv)
   {
       int iv,iq,iss,ise,nmsg,ncpu,ia,id;
       double *sdata,**data;
       double *rdata;
       pickle_v* sbuf;
       int ssize,*iref;
       int nq = 4;
       int nv = 2;
       size_t len;
       uintptr_t ***wptr;

       sdata = new double [nv*nq];
       rdata = new double [nv*nq];
       ssize= nv*nq*sizeof(sdata[0]);
       sbuf= new pickle_v[ssize];
       data = new double *[nv];
       iref = new int [nq];
       subv_cpu(nv,nq,sdata,data);

       nmsg = 2;
       ncpu = 2; 
       wptr = new uintptr_t** [nmsg];
       for(id=0; id<nmsg; id++)
      {
          wptr[id] =new uintptr_t* [ncpu];
          for(ia=0; ia<ncpu; ia++)
         {
             wptr[id][ia] =new uintptr_t [nv];
         }    
      }
       #pragma acc enter data copyin(wptr[0:nmsg])
       for(id=0; id<nmsg; id++)
      {
          #pragma acc enter data copyin(wptr[id][0:ncpu])
          for(ia=0; ia<ncpu; ia++)
         {
             #pragma acc enter data copyin(wptr[id][ia][0:nv])
         }    
      }
       #pragma acc enter data copyin(sdata[0:nv*nq],sbuf[0:ssize],data[0:nv],iref[0:nq])

       #pragma acc parallel loop present(iref[0:nq])
       for(iq=0; iq<nq; iq++) 
      {
          iref[iq] =iq;
      }
       subv(nv,nq,sdata,data);

       #pragma acc parallel loop present(data[0:nv])
       for(iv=0;iv<nv; iv++)
      {
          for(iq=0; iq<nq; iq++)
         {
             data[iv][iq] = iq+iv;
         }   
      }

      //check before pickle
      #pragma acc update self(sdata[0:nv*nq])
       subv_cpu(nv,nq,sdata,data);
       for(iv=0;iv<nv; iv++)
      {
          for(iq=0; iq<nq; iq++)
         {
             cout << iv << " " << iq << " " << data[iv][iq] << "\n";
         }   
      }

       iss = nq/2;
       ise = nq;
       len = 0; 
       for(iv=0; iv<nv; iv++)
      {
          pckle( &len, ise-iss, iref+iss, data[iv], sbuf );
      }

       cout << "------------test unpickle---------\n";
       len = 0;
       subv_cpu(nv,nq,rdata,data);
       for(iv=0;iv<nv; iv++)
      {
          for(iq=0; iq<nq; iq++)
         {
             data[iv][iq] = -1;
         }   
      }
      #pragma acc enter data copyin(rdata[0:nv*nq])
       subv(nv,nq,rdata,data);

      for(id=0; id<nmsg; id++)
     {
         for(ia=0; ia<ncpu; ia++)
        {
            cout << "unpicke message " << id << " in rank " << ia << "\n";
            //#pragma acc kernels present(data[0:nv],wptr[id][ia][0:nv])
            #pragma acc kernels present(data[0:nv],wptr)
            {
                #pragma acc loop seq
                for( iv=0;iv<nv;iv++ )
               {   
                   wptr[id][ia][iv]= (uintptr_t)(data[iv]+iss);
               }
            }

            #pragma acc update self(wptr[id][ia][0:nv])
             len = 0;
             for(iv=0; iv<nv; iv++)
            {
                unpckle( &len, ise-iss, (double*)wptr[id][ia][iv], sbuf );
            }

            //check after unpickle
            #pragma acc update self(rdata[0:nv*nq])
             subv_cpu(nv,nq,rdata,data);
             for(iv=0;iv<nv; iv++)
            {
                for(iq=0; iq<nq; iq++)
               {
                   cout << iv << " " << iq << " " << data[iv][iq] << "\n";
               }      
            }
         }
      }


       #pragma acc exit data delete(sdata,sbuf,data,iref,rdata)
       for(id=0; id<nmsg; id++)
      {
          for(ia=0; ia<ncpu; ia++)
         {
             #pragma acc exit data delete(wptr[id][ia])
         }    
          #pragma acc exit data delete(wptr[id])
      }
       #pragma acc exit data delete(wptr)

       delete[] sdata;
       delete[] rdata;
       delete[] sbuf;
       delete[] data;
       delete[] iref;
       delete[] wptr;
       
       return 0;
   }