#ifndef POTENTIALS_H
#define POTENTIALS_H

#include "precision.hpp"
#include "vec3.h"

#ifdef __CUDACC__
#define CUDA_CALLABLE __host__ __device__
#else
#define CUDA_CALLABLE
#endif

/**
 * Result struct for the Pair Potential
 */
struct ForceAndEnergy {
  real energy;
  Vec3<real> force;

  CUDA_CALLABLE inline static ForceAndEnergy zero() {
    return {0.0, {0.0, 0.0, 0.0}};
  };
};

/**
 * Abstract implementation of a Pair Potential.
 * Pair potentials are potentials which depend solely on the distance
 * between two particles. These do not include multi-body potentials such as
 * EAM
 *
 */
struct PairPotential {
  real m_rcutoffsq;

  PairPotential(real rcutoff) : m_rcutoffsq(rcutoff * rcutoff) {};
#ifdef __CUDACC__
  CUDA_CALLABLE ~PairPotential();
#else
  virtual ~PairPotential() = 0;
#endif

  /**
   * Calculate the force and energy for a specific atom pair based on a
   * displacement vector r.
   */
  CUDA_CALLABLE virtual ForceAndEnergy calc_force_and_energy(Vec3<real> r) = 0;
};

/**
 * Calculate the Lennard-Jones energy and force for the current particle pair
 * described by displacement vector r
 */
struct LennardJones : PairPotential {
  real m_epsilon;
  real m_sigma;

  CUDA_CALLABLE LennardJones(real sigma, real epsilon, real rcutoff)
      : PairPotential(rcutoff), m_epsilon(epsilon), m_sigma(sigma) {};

  CUDA_CALLABLE ForceAndEnergy calc_force_and_energy(Vec3<real> r) {
    real rmagsq = r.squared_norm2();
    if (rmagsq < this->m_rcutoffsq && rmagsq > 0.0) {
      real inv_rmag = 1 / std::sqrt(rmagsq);

      // Pre-Compute the terms (doing this saves on multiple devisions/pow
      // function call)
      real sigma_r = m_sigma * inv_rmag;
      real sigma_r6 = sigma_r * sigma_r * sigma_r * sigma_r * sigma_r * sigma_r;
      real sigma_r12 = sigma_r6 * sigma_r6;

      // Get the energy
      real energy = 4.0 * m_epsilon * (sigma_r12 - sigma_r6);

      // Get the force vector
      real force_mag =
          4.0 * m_epsilon *
          (12.0 * sigma_r12 * inv_rmag - 6.0 * sigma_r6 * inv_rmag);
      Vec3<real> force = r.scale(force_mag * inv_rmag);

      return {energy, force};

    } else {
      return ForceAndEnergy::zero();
    }
  };

  CUDA_CALLABLE ~LennardJones(){};
};

PairPotential::~PairPotential() {};
#endif