typedef float  scalar_t;
typedef float3 vec_t;
typedef float4 data_vec_t;

vec_t get(__global data_vec_t* data,
          unsigned int iParticle)
{
  return data[iParticle].xyz;
}

scalar_t sq(scalar_t x) {
  return pown(x,2);
}

vec_t lennard_jones(__global data_vec_t* pos, unsigned int i, unsigned int j, scalar_t sigma, scalar_t epsilon, vec_t shift, scalar_t cutoff) {
  vec_t iPos = get(pos, i);
  vec_t jPos = get(pos, j) + shift;

  vec_t rij = jPos - iPos;

  scalar_t r = length(rij);
  scalar_t s = pown(sigma / r, 6);
  scalar_t f = 24 * epsilon * s / sq(r) * (1 - 2*s);

  return (r < cutoff) * f * rij;
}

vec_t coulomb(__global data_vec_t* pos, unsigned int i, unsigned int j, scalar_t charge, vec_t shift, scalar_t cutoff) {
  vec_t iPos = get(pos, i);
  vec_t jPos = get(pos, j) + shift;

  vec_t rij = jPos - iPos;
  scalar_t r = length(rij);
  // Coulomb's constant is 138.935458 kJ mol^-1 nm e^-2 in MD units
  scalar_t f = -138.935458 * charge * (1 / pown(r,3));

  return (r < cutoff) * f * rij;
}

vec_t bond(__global data_vec_t* pos, unsigned int i, unsigned int j, scalar_t kb, scalar_t r0) {
  vec_t iPos = get(pos, i);
  vec_t jPos = get(pos, j);

  vec_t rij = jPos - iPos;
  scalar_t r = length(rij);

  return kb * (1 - r0 / r) * rij;
}

vec_t angle(__global data_vec_t* pos, unsigned int iParticle, unsigned int i, unsigned int j, unsigned int k, scalar_t theta0, scalar_t k0) {
  vec_t iPos = get(pos, i);
  vec_t jPos = get(pos, j);
  vec_t kPos = get(pos, k);

  vec_t rij = jPos - iPos;
  vec_t rjk = kPos - jPos;
  vec_t rkj = jPos - kPos;

  scalar_t S = dot(rij, rkj);
  scalar_t D = length(rij) * length(rkj);

  scalar_t theta = acos(S/D);

  scalar_t sd2 = S / pown(D,2);

  scalar_t factor = k0 * sin(theta - theta0) / sin(theta) * (1/D);
  vec_t fi = factor * (sd2 * sq(length(rkj)) * rij - rkj);
  vec_t fk = factor * (sd2 * sq(length(rij)) * rkj - rij);

  return (iParticle == i) *   fi
       + (iParticle == k) *      fk
       + (iParticle == j) * -(fi+fk);
}

vec_t torsion(__global data_vec_t* pos, unsigned int iParticle, unsigned int i, unsigned int j, unsigned int k, unsigned int l, scalar_t phi0, scalar_t k0) {
  vec_t iPos = get(pos, i);
  vec_t jPos = get(pos, j);
  vec_t kPos = get(pos, k);
  vec_t lPos = get(pos, l);

  vec_t rij = jPos - iPos;
  vec_t rjk = kPos - jPos;
  vec_t rik = kPos - iPos;
  vec_t rkl = lPos - kPos;

  vec_t m = cross(rij,rjk);
  vec_t n = cross(rjk,rkl);

  scalar_t phi = M_PI + sign(dot(rij,n)) * acos(clamp(dot(m,n)/(length(m)*length(n)),-1.,1.));
  scalar_t pot = -k0 * sin(phi - phi0);

  vec_t fi = - pot * length(rjk) / sq(length(m)) * m;
  vec_t fl =   pot * length(rjk) / sq(length(n)) * n;
  vec_t fj = -dot(rij,rjk)/sq(length(rjk)) * fi - fi + dot(rkl,rjk)/sq(length(rjk)) * fl;
  vec_t fk = -dot(rkl,rjk)/sq(length(rjk)) * fl - fl + dot(rij,rjk)/sq(length(rjk)) * fi;

  return (iParticle == i) * fi
       + (iParticle == l) * fl
       + (iParticle == j) * fj
       + (iParticle == k) * fk;
}

__kernel void compute_lennard_jones(__global data_vec_t* pos,
                                    __global data_vec_t* force,
                                    __global unsigned int* lj_count,
                                    __global unsigned int* lj_indices,
                                    __global scalar_t* lj_sigma,
                                    __global scalar_t* lj_epsilon,
                                    __global data_vec_t* lj_shift,
                                    scalar_t cutoff)
{
  unsigned int iParticle = get_global_id(0);
  unsigned int nNeighbors = lj_count[iParticle];

  vec_t f = 0;

  for (unsigned int iNeighbor=iParticle*$max_lj; iNeighbor < iParticle*$max_lj + nNeighbors; ++iNeighbor) {
    f += lennard_jones(pos, iParticle,
                       lj_indices[iNeighbor],
                       lj_sigma[iNeighbor],
                       lj_epsilon[iNeighbor],
                       lj_shift[iNeighbor].xyz,
                       cutoff);
  }

  force[iParticle].xyz += f;
}

__kernel void compute_coulomb(__global data_vec_t* pos,
                              __global data_vec_t* force,
                              __global unsigned int* coulomb_count,
                              __global unsigned int* coulomb_indices,
                              __global scalar_t* coulomb_charge,
                              __global data_vec_t* coulomb_shift,
                              scalar_t cutoff)
{
  unsigned int iParticle = get_global_id(0);
  unsigned int nNeighbors = coulomb_count[iParticle];

  vec_t f = 0;

  for (unsigned int iNeighbor=iParticle*$max_coulomb; iNeighbor < iParticle*$max_coulomb + nNeighbors; ++iNeighbor) {
    f += coulomb(pos, iParticle,
                 coulomb_indices[iNeighbor],
                 coulomb_charge[iNeighbor],
                 coulomb_shift[iNeighbor].xyz,
                 cutoff);
  }

  force[iParticle].xyz += f;
}

__kernel void compute_bonds(__global data_vec_t* pos,
                            __global data_vec_t* force,

                            __global unsigned int* bond_count,
                            __global unsigned int* bond_indices,
                            __global scalar_t* bond_kb,
                            __global scalar_t* bond_k0)
{
  unsigned int iParticle = get_global_id(0);

  vec_t f = 0;

  unsigned int nBonds = bond_count[iParticle];
  for (unsigned int iBond=iParticle*$max_bonds; iBond < iParticle*$max_bonds + nBonds; ++iBond) {
    f += bond(pos,
              iParticle,
              bond_indices[iBond],
              bond_kb[iBond],
              bond_k0[iBond]);
  }

  force[iParticle].xyz += f;
}

__kernel void compute_angles(__global data_vec_t* pos,
                             __global data_vec_t* force,

                             __global unsigned int* angle_count,
                             __global unsigned int* angle_indices,
                             __global scalar_t* angle_theta0,
                             __global scalar_t* angle_r0)
{
  unsigned int iParticle = get_global_id(0);

  vec_t f = 0;

  unsigned int nAngles = angle_count[iParticle];
  for (unsigned int iAngle=iParticle*$max_angles; iAngle < iParticle*$max_angles + nAngles; ++iAngle) {
    f += angle(pos, iParticle,
               angle_indices[0*$n_angles+iAngle],
               angle_indices[1*$n_angles+iAngle],
               angle_indices[2*$n_angles+iAngle],
               angle_theta0[iAngle],
               angle_r0[iAngle]);
  }

  force[iParticle].xyz += f;
}

__kernel void compute_torsions(__global data_vec_t* pos,
                               __global data_vec_t* force,

                               __global unsigned int* torsion_count,
                               __global unsigned int* torsion_indices,
                               __global scalar_t* torsion_phi0,
                               __global scalar_t* torsion_r0)
{
  unsigned int iParticle = get_global_id(0);

  vec_t f = 0;

  unsigned int nTorsions = torsion_count[iParticle];
  for (unsigned int iTorsion=iParticle*$max_torsions; iTorsion < iParticle*$max_torsions + nTorsions; ++iTorsion) {
    f += torsion(pos, iParticle,
                 torsion_indices[0*$n_torsions+iTorsion],
                 torsion_indices[1*$n_torsions+iTorsion],
                 torsion_indices[2*$n_torsions+iTorsion],
                 torsion_indices[3*$n_torsions+iTorsion],
                 torsion_phi0[iTorsion],
                 torsion_r0[iTorsion]);
  }

  force[iParticle].xyz += f;
}

__kernel void compute_intramolecular(__global data_vec_t* pos,
                                     __global data_vec_t* force,

                                     __global unsigned int* bond_count,
                                     __global unsigned int* bond_indices,
                                     __global scalar_t* bond_kb,
                                     __global scalar_t* bond_k0,

                                     __global unsigned int* angle_count,
                                     __global unsigned int* angle_indices,
                                     __global scalar_t* angle_theta0,
                                     __global scalar_t* angle_r0,

                                     __global unsigned int* torsion_count,
                                     __global unsigned int* torsion_indices,
                                     __global scalar_t* torsion_phi0,
                                     __global scalar_t* torsion_r0)
{
  unsigned int iParticle = get_global_id(0);

  vec_t f = 0;

  unsigned int nBonds = bond_count[iParticle];
  for (unsigned int iBond=iParticle*$max_bonds; iBond < iParticle*$max_bonds + nBonds; ++iBond) {
    f += bond(pos,
              iParticle,
              bond_indices[iBond],
              bond_kb[iBond],
              bond_k0[iBond]);
  }

  unsigned int nAngles = angle_count[iParticle];
  for (unsigned int iAngle=iParticle*$max_angles; iAngle < iParticle*$max_angles + nAngles; ++iAngle) {
    f += angle(pos, iParticle,
               angle_indices[0*$n_angles+iAngle],
               angle_indices[1*$n_angles+iAngle],
               angle_indices[2*$n_angles+iAngle],
               angle_theta0[iAngle],
               angle_r0[iAngle]);
  }

  unsigned int nTorsions = torsion_count[iParticle];
  for (unsigned int iTorsion=iParticle*$max_torsions; iTorsion < iParticle*$max_torsions + nTorsions; ++iTorsion) {
    f += torsion(pos, iParticle,
                 torsion_indices[0*$n_torsions+iTorsion],
                 torsion_indices[1*$n_torsions+iTorsion],
                 torsion_indices[2*$n_torsions+iTorsion],
                 torsion_indices[3*$n_torsions+iTorsion],
                 torsion_phi0[iTorsion],
                 torsion_r0[iTorsion]);
  }

  force[iParticle].xyz += f;
}

__kernel void evolve_x(__global data_vec_t* pos,
                       __global data_vec_t* vel,
                       __global data_vec_t* force_prev,
                       __global data_vec_t* force_curr,
                       scalar_t tau)
{
  unsigned int iParticle = get_global_id(0);

  vec_t f_prev = get(force_prev, iParticle);
  vec_t f_curr = get(force_curr, iParticle);

  vec_t p = get(pos, iParticle);
  vec_t v = get(vel, iParticle);

  scalar_t a = tau * 0.5 / pos[iParticle].w;
  p += tau * (v + a * f_curr);

  pos[iParticle].xyz = p;
  force_prev[iParticle].xyz = f_curr;
  force_curr[iParticle].xyz = 0;
}

__kernel void evolve_v(__global data_vec_t* pos,
                       __global data_vec_t* vel,
                       __global data_vec_t* force_prev,
                       __global data_vec_t* force_curr,
                       scalar_t tau)
{
  unsigned int iParticle = get_global_id(0);

  vec_t f_prev = get(force_prev, iParticle);
  vec_t f_curr = get(force_curr, iParticle);

  vec_t v = get(vel, iParticle);

  scalar_t a = tau * 0.5 / pos[iParticle].w;
  v += a * (f_prev + f_curr);

  vel[iParticle].xyz = v;
}

scalar_t lennard_jones_sigma(scalar_t iMass)
{
  return $lj_sigma_expr;
}

scalar_t lennard_jones_epsilon(scalar_t iMass)
{
  return $lj_epsilon_expr;
}

scalar_t coulomb_q(scalar_t iMass)
{
  return $coulomb_charge_expr;
}

__kernel void update_potential_neighborhoods(__global data_vec_t* pos,
                                             __global unsigned int* molecules,

                                             __global unsigned int* lj_count,
                                             __global unsigned int* lj_indices,
                                             __global scalar_t*     lj_sigma,
                                             __global scalar_t*     lj_epsilon,
                                             __global data_vec_t*   lj_shift,

                                             __global unsigned int* coulomb_count,
                                             __global unsigned int* coulomb_indices,
                                             __global scalar_t*     coulomb_charge,
                                             __global data_vec_t*   coulomb_shift,

                                             scalar_t cutoff,
                                             scalar_t skin)
{
  unsigned int iParticle = get_global_id(0);

  __global unsigned int* lj_i_count   = lj_count   + $max_lj * iParticle;
  __global unsigned int* lj_i_indices = lj_indices + $max_lj * iParticle;
  __global scalar_t*     lj_i_sigma   = lj_sigma   + $max_lj * iParticle;
  __global scalar_t*     lj_i_epsilon = lj_epsilon + $max_lj * iParticle;
  __global data_vec_t*   lj_i_shift   = lj_shift   + $max_lj * iParticle;

  __global unsigned int* coulomb_i_count   = coulomb_count   + $max_coulomb * iParticle;
  __global unsigned int* coulomb_i_indices = coulomb_indices + $max_coulomb * iParticle;
  __global scalar_t*     coulomb_i_charge  = coulomb_charge  + $max_coulomb * iParticle;
  __global data_vec_t*   coulomb_i_shift   = coulomb_shift   + $max_coulomb * iParticle;

  for (unsigned int i=0; i < $max_lj; ++i) {
    lj_i_indices[i] = 0;
    lj_i_sigma[i]   = 0;
    lj_i_epsilon[i] = 0;
    lj_i_shift[i] = 0;

    coulomb_i_indices[i] = 0;
    coulomb_i_charge[i]  = 0;
    coulomb_i_shift[i] = 0;
  }

  unsigned int idx = 0;

  vec_t iPos = get(pos, iParticle);

  for (unsigned int jParticle=0; jParticle < $n_atoms; ++jParticle) {
    for (int sX=-1; sX <= 1; ++sX) {
      for (int sY=-1; sY <= 1; ++sY) {
        for (int sZ=-1; sZ <= 1; ++sZ) {
          vec_t shift = $domain_size * (vec_t)(sX,sY,sZ);
          vec_t jPos = get(pos, jParticle) + shift;

          scalar_t r = length(jPos - iPos);

          unsigned int iMolecule = molecules[iParticle];
          unsigned int jMolecule = molecules[jParticle];

          scalar_t iMass = pos[iParticle].w;
          scalar_t jMass = pos[jParticle].w;

          scalar_t mask = (r < cutoff + skin) * (iParticle != jParticle) * (iMolecule != jMolecule);

          lj_i_indices[idx]   += mask * jParticle;
          lj_i_sigma[idx]     += mask * 0.5*(lennard_jones_sigma(iMass) + lennard_jones_sigma(jMass));
          lj_i_epsilon[idx]   += mask * sqrt(lennard_jones_epsilon(iMass) * lennard_jones_epsilon(jMass));
          lj_i_shift[idx].xyz += mask * shift;

          coulomb_i_indices[idx]   += mask * jParticle;
          coulomb_i_charge[idx]    += mask * coulomb_q(iMass) * coulomb_q(jMass);
          coulomb_i_shift[idx].xyz += mask * shift;

          idx += mask;
        }
      }
    }
  }

  lj_count[iParticle] = idx;
  coulomb_count[iParticle] = idx;
}

__kernel void wrap_molecules(__global data_vec_t* pos,
                             __global unsigned int* molecules)
{
  unsigned int iMolecule = get_global_id(0);

  int nAtoms = 0;
  int xShift = 0;
  int yShift = 0;
  int zShift = 0;

  for (unsigned int jParticle=0; jParticle < $n_atoms; ++jParticle) {
    nAtoms += (iMolecule == molecules[jParticle]);
    xShift -= (iMolecule == molecules[jParticle]) * (pos[jParticle].x > $domain_size);
    xShift += (iMolecule == molecules[jParticle]) * (pos[jParticle].x < 0.);
    yShift -= (iMolecule == molecules[jParticle]) * (pos[jParticle].y > $domain_size);
    yShift += (iMolecule == molecules[jParticle]) * (pos[jParticle].y < 0.);
    zShift -= (iMolecule == molecules[jParticle]) * (pos[jParticle].z > $domain_size);
    zShift += (iMolecule == molecules[jParticle]) * (pos[jParticle].z < 0.);
  }

  for (unsigned int jParticle=0; jParticle < $n_atoms; ++jParticle) {
    if (molecules[jParticle] == iMolecule) {
      pos[jParticle].x += (xShift ==  nAtoms) * $domain_size;
      pos[jParticle].x -= (xShift == -nAtoms) * $domain_size;
      pos[jParticle].y += (yShift ==  nAtoms) * $domain_size;
      pos[jParticle].y -= (yShift == -nAtoms) * $domain_size;
      pos[jParticle].z += (zShift ==  nAtoms) * $domain_size;
      pos[jParticle].z -= (zShift == -nAtoms) * $domain_size;
    }
  }
}

__kernel void update_view_links(__global data_vec_t* pos,
                                __global unsigned int* link_indices,
                                __global data_vec_t* link_pos)
{
  unsigned int iLink = get_global_id(0);

  link_pos[2*iLink+0].xyz = pos[link_indices[2*iLink+0]].xyz;
  link_pos[2*iLink+1].xyz = pos[link_indices[2*iLink+1]].xyz;
}