/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator 

   Original Version:
   http://lammps.sandia.gov, Sandia National Laboratories
   Steve Plimpton, sjplimp@sandia.gov 

   See the README file in the top-level LAMMPS directory. 

   ----------------------------------------------------------------------- 

   USER-CUDA Package and associated modifications:
   https://sourceforge.net/projects/lammpscuda/ 

   Christian Trott, christian.trott@tu-ilmenau.de
   Lars Winterfeld, lars.winterfeld@tu-ilmenau.de
   Theoretical Physics II, University of Technology Ilmenau, Germany 

   See the README file in the USER-CUDA directory. 

   This software is distributed under the GNU General Public License.
------------------------------------------------------------------------- */
#define Pi F_F(3.1415926535897932384626433832795)
#define PI Pi
#define PI2 F_F(0.5)*Pi
#define PI4 F_F(0.25)*Pi



__device__ void twobody(int iparam, F_FLOAT rsq, F_FLOAT &fforce,
          int eflag, ENERGY_FLOAT &eng)
{
  F_FLOAT r,rp,rq,rainv,expsrainv;

  r = sqrt(rsq);
  rp = pow(r,-params_sw[iparam].powerp);
  rq = pow(r,-params_sw[iparam].powerq);
  rainv = 1.0 / (r - params_sw[iparam].cut);
  expsrainv = exp(params_sw[iparam].sigma * rainv);
  fforce = (params_sw[iparam].c1*rp - params_sw[iparam].c2*rq +
      (params_sw[iparam].c3*rp -params_sw[iparam].c4*rq) * rainv*rainv*r) * expsrainv / rsq;
  if (eflag) eng += (params_sw[iparam].c5*rp - params_sw[iparam].c6*rq) * expsrainv;
}

__device__ void threebody(int paramij, int paramik, int paramijk,
           F_FLOAT4& delr1,
           F_FLOAT4& delr2,
           F_FLOAT3& fj, F_FLOAT3& fk, int eflag,ENERGY_FLOAT &eng)
{
  F_FLOAT r1,rinvsq1,rainv1,gsrainv1,gsrainvsq1,expgsrainv1;
  F_FLOAT r2,rinvsq2,rainv2,gsrainv2,gsrainvsq2,expgsrainv2;
  F_FLOAT rinv12,cs,delcs,delcssq,facexp,facrad,frad1,frad2;
  F_FLOAT facang,facang12,csfacang,csfac1,csfac2;

  r1 = sqrt(delr1.w);
  rinvsq1 = F_F(1.0)/delr1.w;
  rainv1 = F_F(1.0)/(r1 - params_sw[paramij].cut);
  gsrainv1 = params_sw[paramij].sigma_gamma * rainv1;
  gsrainvsq1 = gsrainv1*rainv1/r1;
  expgsrainv1 = exp(gsrainv1);

  r2 = sqrt(delr2.w);
  rinvsq2 = F_F(1.0)/delr2.w;
  rainv2 = F_F(1.0)/(r2 - params_sw[paramik].cut);
  gsrainv2 = params_sw[paramik].sigma_gamma * rainv2;
  gsrainvsq2 = gsrainv2*rainv2/r2;
  expgsrainv2 = exp(gsrainv2);

  rinv12 = F_F(1.0)/(r1*r2);
  cs = (delr1.x*delr2.x + delr1.y*delr2.y + delr1.z*delr2.z) * rinv12;
  delcs = cs - params_sw[paramijk].costheta;
  delcssq = delcs*delcs;

  facexp = expgsrainv1*expgsrainv2;

  // facrad = sqrt(paramij->lambda_epsilon*paramik->lambda_epsilon) *
  //          facexp*delcssq;

  facrad = params_sw[paramijk].lambda_epsilon * facexp*delcssq;
  frad1 = facrad*gsrainvsq1;
  frad2 = facrad*gsrainvsq2;
  facang = params_sw[paramijk].lambda_epsilon2 * facexp*delcs;
  facang12 = rinv12*facang;
  csfacang = cs*facang;
  csfac1 = rinvsq1*csfacang;

  fj.x = delr1.x*(frad1+csfac1)-delr2.x*facang12;
  fj.y = delr1.y*(frad1+csfac1)-delr2.y*facang12;
  fj.z = delr1.z*(frad1+csfac1)-delr2.z*facang12;

  csfac2 = rinvsq2*csfacang;

  fk.x = delr2.x*(frad2+csfac2)-delr1.x*facang12;
  fk.y = delr2.y*(frad2+csfac2)-delr1.y*facang12;
  fk.z = delr2.z*(frad2+csfac2)-delr1.z*facang12;

  if (eflag) eng+= F_F(2.0)*facrad;
}

__device__ void threebody_fj(int paramij, int paramik, int paramijk,
           F_FLOAT4& delr1,
           F_FLOAT4& delr2,
           F_FLOAT3& fj)
{
  F_FLOAT r1,rinvsq1,rainv1,gsrainv1,gsrainvsq1,expgsrainv1;
  F_FLOAT r2,rainv2,gsrainv2,expgsrainv2;
  F_FLOAT rinv12,cs,delcs,delcssq,facexp,facrad,frad1;
  F_FLOAT facang,facang12,csfacang,csfac1;

  r1 = sqrt(delr1.w);
  rinvsq1 = F_F(1.0)/delr1.w;
  rainv1 = F_F(1.0)/(r1 - params_sw[paramij].cut);
  gsrainv1 = params_sw[paramij].sigma_gamma * rainv1;
  gsrainvsq1 = gsrainv1*rainv1/r1;
  expgsrainv1 = exp(gsrainv1);

  r2 = sqrt(delr2.w);
  rainv2 = F_F(1.0)/(r2 - params_sw[paramik].cut);
  gsrainv2 = params_sw[paramik].sigma_gamma * rainv2;
  expgsrainv2 = exp(gsrainv2);

  rinv12 = F_F(1.0)/(r1*r2);
  cs = (delr1.x*delr2.x + delr1.y*delr2.y + delr1.z*delr2.z) * rinv12;
  delcs = cs - params_sw[paramijk].costheta;
  delcssq = delcs*delcs;

  facexp = expgsrainv1*expgsrainv2;

  // facrad = sqrt(paramij->lambda_epsilon*paramik->lambda_epsilon) *
  //          facexp*delcssq;

  facrad = params_sw[paramijk].lambda_epsilon * facexp*delcssq;
  frad1 = facrad*gsrainvsq1;
  facang = params_sw[paramijk].lambda_epsilon2 * facexp*delcs;
  facang12 = rinv12*facang;
  csfacang = cs*facang;
  csfac1 = rinvsq1*csfacang;

  fj.x = delr1.x*(frad1+csfac1)-delr2.x*facang12;
  fj.y = delr1.y*(frad1+csfac1)-delr2.y*facang12;
  fj.z = delr1.z*(frad1+csfac1)-delr2.z*facang12;
}


__global__ void Pair_SW_Kernel_TpA_RIJ()//F_FLOAT4* _glob_r_ij,int* _glob_numneigh_red,int* _glob_neighbors_red,int* _glob_neightype_red)
{
  int ii = (blockIdx.x*gridDim.y+blockIdx.y)*blockDim.x+threadIdx.x;
  if( ii >= _nall ) return;

  X_FLOAT4 myxtype;
  F_FLOAT4 delij;
  F_FLOAT xtmp,ytmp,ztmp;
  int itype,jnum,i,j;
  int* jlist;
  int neigh_red = 0;
  i = ii;//_ilist[ii];
  myxtype = fetchXType(i);

  xtmp=myxtype.x;
  ytmp=myxtype.y;
  ztmp=myxtype.z;
  itype=map[(static_cast <int> (myxtype.w))];

  jnum = _numneigh[i];
  jlist = &_neighbors[i];

  __syncthreads();
  for (int jj = 0; jj < jnum; jj++)
  {
    if(jj<jnum)
    {

      j = jlist[jj*_nall];
      j &= NEIGHMASK;
      myxtype = fetchXType(j);
      delij.x = xtmp - myxtype.x;
      delij.y = ytmp - myxtype.y;
      delij.z = ztmp - myxtype.z;
      int jtype = map[(static_cast <int> (myxtype.w))];
      int iparam_ij = elem2param[(itype*nelements+jtype)*nelements+jtype];
      delij.w = vec3_dot(delij,delij);
      if (delij.w < params_sw[iparam_ij].cutsq)
      {
        _glob_neighbors_red[i+neigh_red*_nall]=j;
        _glob_neightype_red[i+neigh_red*_nall]=jtype;
        _glob_r_ij[i+neigh_red*_nall]=delij;
        neigh_red++;
      }
    }
  }
  _glob_numneigh_red[i]=neigh_red;
}


 template <int eflag, int vflagm>
 __global__ void Pair_SW_Kernel_TpA(int eflag_atom,int vflag_atom)//,F_FLOAT* _glob_zeta_ij,F_FLOAT4* _glob_r_ij,int* _glob_numneigh_red,int* _glob_neighbors_red,int* _glob_neightype_red)
{
  ENERGY_FLOAT evdwl = ENERGY_F(0.0);

  ENERGY_FLOAT* sharedE = &sharedmem[threadIdx.x];
  ENERGY_FLOAT* sharedV = &sharedmem[threadIdx.x];

  F_FLOAT* shared_F_F = (F_FLOAT*) sharedmem;
  if((eflag||eflag_atom)&&(vflagm||vflag_atom)) shared_F_F = (F_FLOAT*) &sharedmem[7*blockDim.x];
  else
  if(eflag) shared_F_F = (F_FLOAT*) &sharedmem[blockDim.x];
  else
  if(vflagm) shared_F_F = (F_FLOAT*) &sharedmem[6*blockDim.x];
  shared_F_F+=threadIdx.x;

  if(eflag_atom||eflag)
  {
    sharedE[0] = ENERGY_F(0.0);
    sharedV += blockDim.x;
  }

  if(vflagm||vflag_atom)
  {
    sharedV[0*blockDim.x] = ENERGY_F(0.0);
    sharedV[1*blockDim.x] = ENERGY_F(0.0);
    sharedV[2*blockDim.x] = ENERGY_F(0.0);
    sharedV[3*blockDim.x] = ENERGY_F(0.0);
    sharedV[4*blockDim.x] = ENERGY_F(0.0);
    sharedV[5*blockDim.x] = ENERGY_F(0.0);
  }

  int jnum_red=0;
#define fxtmp shared_F_F[0]
#define fytmp shared_F_F[blockDim.x]
#define fztmp shared_F_F[2*blockDim.x]
//#define jnum_red (static_cast <int> (shared_F_F[3*blockDim.x]))

  int ii = (blockIdx.x*gridDim.y+blockIdx.y)*blockDim.x+threadIdx.x;
  X_FLOAT4 myxtype_i,myxtype_j,myxtype_k;
  F_FLOAT4 delij,delik,deljk;
  F_FLOAT fpair;

  int itype,i,j;
  int* jlist_red;

  if(ii < _inum)
  {
    i = _ilist[ii];

    if(vflagm)
    myxtype_i=fetchXType(i);
    //itype=map[(static_cast <int> (myxtype_i.w))];
    itype=map[_type[i]];


    fxtmp = F_F(0.0);
    fytmp = F_F(0.0);
    fztmp = F_F(0.0);


    //shared_F_F[3*blockDim.x] = _glob_numneigh_red[i];
    jnum_red = _glob_numneigh_red[i];
    jlist_red = &_glob_neighbors_red[i];
  }
  __syncthreads();
#pragma unroll 1
  for (int jj = 0; jj < jnum_red; jj++)
  {
    if(i < _nlocal)
    {
      fpair=F_F(0.0);
      j = jlist_red[jj*_nall];
      j &= NEIGHMASK;

      if(vflagm)
      myxtype_j = fetchXType(j);

      int jtype = _glob_neightype_red[i+jj*_nall];
      delij = _glob_r_ij[i+jj*_nall];

      volatile int iparam_ij = elem2param[(itype*nelements+jtype)*nelements+jtype];
      volatile int iparam_ji = elem2param[(jtype*nelements+itype)*nelements+itype];

      if (delij.w<params_sw[iparam_ij].cutsq)
      {
        F_FLOAT dxfp,dyfp,dzfp;
        twobody(iparam_ij,delij.w,fpair,eflag,evdwl);
        fxtmp += dxfp = delij.x*fpair;
        fytmp += dyfp = delij.y*fpair;
        fztmp += dzfp = delij.z*fpair;
        if(vflagm)
        {
          sharedV[0 * blockDim.x]+= delij.x*dxfp;
          sharedV[1 * blockDim.x]+= delij.y*dyfp;
          sharedV[2 * blockDim.x]+= delij.z*dzfp;
          sharedV[3 * blockDim.x]+= delij.x*dyfp;
          sharedV[4 * blockDim.x]+= delij.x*dzfp;
          sharedV[5 * blockDim.x]+= delij.y*dzfp;
        }






        vec3_scale(F_F(-1.0),delij,delij);

#pragma unroll 1
        for (int kk = jj+1; kk < jnum_red; kk++) {
          int k = jlist_red[kk*_nall];
          k &= NEIGHMASK;

          if(vflagm)
          myxtype_k = fetchXType(k);

          delik = _glob_r_ij[i+kk*_nall];

          int ktype = _glob_neightype_red[i+kk*_nall];
          int iparam_ik = elem2param[(itype*nelements+ktype)*nelements+ktype];
          int iparam_ijk = elem2param[(itype*nelements+jtype)*nelements+ktype];
          vec3_scale(F_F(-1.0),delik,delik);
          if (delik.w <= params_sw[iparam_ijk].cutsq)
          {
            F_FLOAT3 fj,fk;
            threebody(iparam_ij,iparam_ik,iparam_ijk,
              delij,delik,fj,fk,eflag,evdwl);
            fxtmp -= fj.x + fk.x;
            fytmp -= fj.y + fk.y;
            fztmp -= fj.z + fk.z;

            if(vflagm)
            {
              sharedV[0 * blockDim.x]-= ENERGY_F(2.0)*myxtype_i.x*(fj.x+fk.x);
              sharedV[1 * blockDim.x]-= ENERGY_F(2.0)*myxtype_i.y*(fj.y+fk.y);
              sharedV[2 * blockDim.x]-= ENERGY_F(2.0)*myxtype_i.z*(fj.z+fk.z);
              sharedV[3 * blockDim.x]-= ENERGY_F(2.0)*myxtype_i.x*(fj.y+fk.y);
              sharedV[4 * blockDim.x]-= ENERGY_F(2.0)*myxtype_i.x*(fj.z+fk.z);
              sharedV[5 * blockDim.x]-= ENERGY_F(2.0)*myxtype_i.y*(fj.z+fk.z);

              sharedV[0 * blockDim.x]+= ENERGY_F(2.0)*myxtype_j.x*fj.x;
              sharedV[1 * blockDim.x]+= ENERGY_F(2.0)*myxtype_j.y*fj.y;
              sharedV[2 * blockDim.x]+= ENERGY_F(2.0)*myxtype_j.z*fj.z;
              sharedV[3 * blockDim.x]+= ENERGY_F(2.0)*myxtype_j.x*fj.y;
              sharedV[4 * blockDim.x]+= ENERGY_F(2.0)*myxtype_j.x*fj.z;
              sharedV[5 * blockDim.x]+= ENERGY_F(2.0)*myxtype_j.y*fj.z;

              sharedV[0 * blockDim.x]+= ENERGY_F(2.0)*myxtype_k.x*fk.x;
              sharedV[1 * blockDim.x]+= ENERGY_F(2.0)*myxtype_k.y*fk.y;
              sharedV[2 * blockDim.x]+= ENERGY_F(2.0)*myxtype_k.z*fk.z;
              sharedV[3 * blockDim.x]+= ENERGY_F(2.0)*myxtype_k.x*fk.y;
              sharedV[4 * blockDim.x]+= ENERGY_F(2.0)*myxtype_k.x*fk.z;
              sharedV[5 * blockDim.x]+= ENERGY_F(2.0)*myxtype_k.y*fk.z;
            }
          }
        }

        int j_jnum_red = _glob_numneigh_red[j];
        int* j_jlist_red = &_glob_neighbors_red[j];

        int j_ii = 0;
//#pragma unroll 1
        for (int j_kk = 0; j_kk < j_jnum_red; j_kk++) {
          if(j_jlist_red[j_kk*_nall]==i) j_ii = j_kk;
        }

#pragma unroll 1
        for (int kk = 0; kk < j_jnum_red; kk++)
        {
          if (j_ii == kk) continue;
          int k = j_jlist_red[kk*_nall];
          k &= NEIGHMASK;
          deljk = _glob_r_ij[j+kk*_nall];
          vec3_scale(F_F(-1.0),deljk,deljk);
          int ktype = _glob_neightype_red[j+kk*_nall];

          int iparam_ji = elem2param[(jtype*nelements+itype)*nelements+itype];
          int iparam_jk = elem2param[(jtype*nelements+ktype)*nelements+ktype];
          int iparam_jik = elem2param[(jtype*nelements+itype)*nelements+ktype];


          vec3_scale(F_F(-1.0),delij,delij);
          if (deljk.w <= params_sw[iparam_jik].cutsq)
          {
            F_FLOAT3 fj;

            threebody_fj(iparam_ji,iparam_jk,iparam_jik,
              delij,deljk,fj);
            fxtmp += fj.x;
            fytmp += fj.y;
            fztmp += fj.z;

          }
          vec3_scale(F_F(-1.0),delij,delij);
        }
      }
    }

  }
  __syncthreads();
  if(ii < _inum)
  {
      F_FLOAT* my_f;
      if(_collect_forces_later)
      {
        ENERGY_FLOAT* buffer = (ENERGY_FLOAT*) _buffer;
        if(eflag)
        {
          buffer=&buffer[1 * gridDim.x * gridDim.y];
        }
        if(vflagm)
        {
          buffer=&buffer[6 * gridDim.x * gridDim.y];
        }
        my_f = (F_FLOAT*) buffer;
        my_f += i;
        *my_f = fxtmp; my_f += _nmax;
        *my_f = fytmp; my_f += _nmax;
        *my_f = fztmp;
      }
      else
      {
        my_f = _f + i;
        *my_f += fxtmp; my_f += _nmax;
        *my_f += fytmp; my_f += _nmax;
        *my_f += fztmp;
      }
   }
   __syncthreads();

   if(eflag)
   {
     sharedE[0] = evdwl;
   }
   if(eflag_atom && i<_nlocal)
   {
     _eatom[i] = ENERGY_F(0.5) * evdwl;
   }

   if(vflag_atom && i<_nlocal)
   {
     _vatom[i]         = ENERGY_F(0.5) * sharedV[0 * blockDim.x];
     _vatom[i+_nmax]   = ENERGY_F(0.5) * sharedV[1 * blockDim.x];
     _vatom[i+2*_nmax] = ENERGY_F(0.5) * sharedV[2 * blockDim.x];
     _vatom[i+3*_nmax] = ENERGY_F(0.5) * sharedV[3 * blockDim.x];
     _vatom[i+4*_nmax] = ENERGY_F(0.5) * sharedV[4 * blockDim.x];
     _vatom[i+5*_nmax] = ENERGY_F(0.5) * sharedV[5 * blockDim.x];
   }
   if(vflagm&&eflag) PairVirialCompute_A_Kernel_Template<1,1>();
   else if(eflag) PairVirialCompute_A_Kernel_Template<1,0>();
   else if(vflagm) PairVirialCompute_A_Kernel_Template<0,1>();
#undef fxtmp
#undef fytmp
#undef fztmp
//#undef jnum_red
}