#include "pair_potentials.cuh"
#include "precision.hpp"
#include "gtest/gtest.h"
#include <cmath>

class LennardJonesTest : public ::testing::Test {
protected:
  void SetUp() override {
    // Default parameters
    sigma = 1.0;
    epsilon = 1.0;
    r_cutoff = 2.5;

    // Create default LennardJones object
    lj = new LennardJones(sigma, epsilon, r_cutoff);
  }

  void TearDown() override { delete lj; }

  real sigma;
  real epsilon;
  real r_cutoff;
  LennardJones *lj;

  // Helper function to compare Vec3 values with tolerance
  void expect_vec3_near(const Vec3<real> &expected, const Vec3<real> &actual,
                        real tolerance) {
    EXPECT_NEAR(expected.x, actual.x, tolerance);
    EXPECT_NEAR(expected.y, actual.y, tolerance);
    EXPECT_NEAR(expected.z, actual.z, tolerance);
  }
};

TEST_F(LennardJonesTest, ZeroDistance) {
  // At zero distance, the calculation should return zero force and energy
  Vec3<real> r = {0.0, 0.0, 0.0};
  auto result = lj->calc_force_and_energy(r);

  EXPECT_EQ(0.0, result.energy);
  expect_vec3_near({0.0, 0.0, 0.0}, result.force, 1e-10);
}

TEST_F(LennardJonesTest, BeyondCutoff) {
  // Distance beyond cutoff should return zero force and energy
  Vec3<real> r = {3.0, 0.0, 0.0}; // 3.0 > r_cutoff (2.5)
  auto result = lj->calc_force_and_energy(r);

  EXPECT_EQ(0.0, result.energy);
  expect_vec3_near({0.0, 0.0, 0.0}, result.force, 1e-10);
}

TEST_F(LennardJonesTest, AtMinimum) {
  // The LJ potential has a minimum at r = 2^(1/6) * sigma
  real min_dist = std::pow(2.0, 1.0 / 6.0) * sigma;
  Vec3<real> r = {min_dist, 0.0, 0.0};
  auto result = lj->calc_force_and_energy(r);

  // At minimum, force should be close to zero
  EXPECT_NEAR(-epsilon, result.energy, 1e-10);
  expect_vec3_near({0.0, 0.0, 0.0}, result.force, 1e-10);
}

TEST_F(LennardJonesTest, AtEquilibrium) {
  // At r = sigma, the energy should be zero and force should be repulsive
  Vec3<real> r = {sigma, 0.0, 0.0};
  auto result = lj->calc_force_and_energy(r);

  EXPECT_NEAR(0.0, result.energy, 1e-10);
  EXPECT_GT(result.force.x,
            0.0); // Force should be repulsive (positive x-direction)
  EXPECT_NEAR(0.0, result.force.y, 1e-10);
  EXPECT_NEAR(0.0, result.force.z, 1e-10);
}

TEST_F(LennardJonesTest, RepulsiveRegion) {
  // Test in the repulsive region (r < sigma)
  Vec3<real> r = {0.8 * sigma, 0.0, 0.0};
  auto result = lj->calc_force_and_energy(r);

  // Energy should be positive and force should be repulsive
  EXPECT_GT(result.energy, 0.0);
  EXPECT_GT(result.force.x, 0.0); // Force should be repulsive
}

TEST_F(LennardJonesTest, AttractiveRegion) {
  // Test in the attractive region (sigma < r < r_min)
  Vec3<real> r = {1.5 * sigma, 0.0, 0.0};
  auto result = lj->calc_force_and_energy(r);

  // Energy should be negative and force should be attractive
  EXPECT_LT(result.energy, 0.0);
  EXPECT_LT(result.force.x,
            0.0); // Force should be attractive (negative x-direction)
}

TEST_F(LennardJonesTest, ArbitraryDirection) {
  // Test with a vector in an arbitrary direction
  Vec3<real> r = {1.0, 1.0, 1.0};
  auto result = lj->calc_force_and_energy(r);

  // The force should be in the same direction as r but opposite sign
  // (attractive region)
  real r_mag = std::sqrt(r.squared_norm2());

  // Calculate expected force direction (should be along -r)
  Vec3<real> normalized_r = r.scale(1.0 / r_mag);
  real force_dot_r = result.force.x * normalized_r.x +
                     result.force.y * normalized_r.y +
                     result.force.z * normalized_r.z;

  // In this case, we're at r = sqrt(3) * sigma which is in attractive region
  EXPECT_LT(force_dot_r, 0.0); // Force should be attractive

  // Force should be symmetric in all dimensions for this vector
  EXPECT_NEAR(result.force.x, result.force.y, 1e-10);
  EXPECT_NEAR(result.force.y, result.force.z, 1e-10);
}

TEST_F(LennardJonesTest, ParameterVariation) {
  // Test with different parameter values
  real new_sigma = 2.0;
  real new_epsilon = 0.5;
  real new_r_cutoff = 5.0;

  LennardJones lj2(new_sigma, new_epsilon, new_r_cutoff);

  Vec3<real> r = {2.0, 0.0, 0.0};
  auto result1 = lj->calc_force_and_energy(r);
  auto result2 = lj2.calc_force_and_energy(r);

  // Results should be different with different parameters
  EXPECT_NE(result1.energy, result2.energy);
  EXPECT_NE(result1.force.x, result2.force.x);
}

TEST_F(LennardJonesTest, ExactValueCheck) {
  // Test with pre-calculated values for a specific case
  LennardJones lj_exact(1.0, 1.0, 3.0);
  Vec3<real> r = {1.5, 0.0, 0.0};
  auto result = lj_exact.calc_force_and_energy(r);

  // Pre-calculated values (you may need to adjust these based on your specific
  // implementation)
  real expected_energy =
      4.0 * (std::pow(1.0 / 1.5, 12) - std::pow(1.0 / 1.5, 6));
  real expected_force =
      24.0 * (std::pow(1.0 / 1.5, 6) - 2.0 * std::pow(1.0 / 1.5, 12)) / 1.5;

  EXPECT_NEAR(expected_energy, result.energy, 1e-10);
  EXPECT_NEAR(-expected_force, result.force.x,
              1e-10); // Negative because force is attractive
  EXPECT_NEAR(0.0, result.force.y, 1e-10);
  EXPECT_NEAR(0.0, result.force.z, 1e-10);
}

TEST_F(LennardJonesTest, NearCutoff) {
  // Test behavior just inside and just outside the cutoff
  real inside_cutoff = r_cutoff - 0.01;
  real outside_cutoff = r_cutoff + 0.01;

  Vec3<real> r_inside = {inside_cutoff, 0.0, 0.0};
  Vec3<real> r_outside = {outside_cutoff, 0.0, 0.0};

  auto result_inside = lj->calc_force_and_energy(r_inside);
  auto result_outside = lj->calc_force_and_energy(r_outside);

  // Inside should have non-zero values
  EXPECT_NE(0.0, result_inside.energy);
  EXPECT_NE(0.0, result_inside.force.x);

  // Outside should be zero
  EXPECT_EQ(0.0, result_outside.energy);
  expect_vec3_near({0.0, 0.0, 0.0}, result_outside.force, 1e-10);
}