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

// Structure to hold test results from device
struct TestResults {
  bool zero_distance_pass;
  bool beyond_cutoff_pass;
  bool at_minimum_pass;
  bool at_equilibrium_pass;
  bool repulsive_region_pass;
  bool attractive_region_pass;
  bool arbitrary_direction_pass;
  bool parameter_variation_pass;
  bool exact_value_check_pass;
  bool near_cutoff_pass;

  // Additional result data for exact checks
  real energy_values[10];
  Vec3<real> force_values[10];
};

// Check if two Vec3 values are close within tolerance
__device__ bool vec3_near(const Vec3<real> &a, const Vec3<real> &b,
                          real tolerance) {
  return (fabs(a.x - b.x) < tolerance) && (fabs(a.y - b.y) < tolerance) &&
         (fabs(a.z - b.z) < tolerance);
}

// Device kernel to run all tests
__global__ void lennard_jones_test_kernel(TestResults *results) {
  // Default parameters
  real sigma = 1.0;
  real epsilon = 1.0;
  real r_cutoff = 2.5;
  real tolerance = 1e-10;

  // Create LennardJones object on device
  LennardJones lj(sigma, epsilon, r_cutoff);

  // Zero Distance Test
  {
    Vec3<real> r = {0.0, 0.0, 0.0};
    auto result = lj.calc_force_and_energy(r);
    results->energy_values[0] = result.energy;
    results->force_values[0] = result.force;
    results->zero_distance_pass =
        (result.energy == 0.0) &&
        vec3_near(Vec3<real>{0.0, 0.0, 0.0}, result.force, tolerance);
  }

  // Beyond Cutoff Test
  {
    Vec3<real> r = {3.0, 0.0, 0.0};
    auto result = lj.calc_force_and_energy(r);
    results->energy_values[1] = result.energy;
    results->force_values[1] = result.force;
    results->beyond_cutoff_pass =
        (result.energy == 0.0) &&
        vec3_near(Vec3<real>{0.0, 0.0, 0.0}, result.force, tolerance);
  }

  // At Minimum Test
  {
    real min_dist = 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);
    results->energy_values[2] = result.energy;
    results->force_values[2] = result.force;
    results->at_minimum_pass =
        (fabs(result.energy + epsilon) < tolerance) &&
        vec3_near(Vec3<real>{0.0, 0.0, 0.0}, result.force, tolerance);
  }

  // At Equilibrium Test
  {
    Vec3<real> r = {sigma, 0.0, 0.0};
    auto result = lj.calc_force_and_energy(r);
    results->energy_values[3] = result.energy;
    results->force_values[3] = result.force;
    results->at_equilibrium_pass = (fabs(result.energy) < tolerance) &&
                                   (result.force.x > 0.0) &&
                                   (fabs(result.force.y) < tolerance) &&
                                   (fabs(result.force.z) < tolerance);
  }

  // Repulsive Region Test
  {
    Vec3<real> r = {0.8 * sigma, 0.0, 0.0};
    auto result = lj.calc_force_and_energy(r);
    results->energy_values[4] = result.energy;
    results->force_values[4] = result.force;
    results->repulsive_region_pass =
        (result.energy > 0.0) && (result.force.x > 0.0);
  }

  // Attractive Region Test
  {
    Vec3<real> r = {1.5 * sigma, 0.0, 0.0};
    auto result = lj.calc_force_and_energy(r);
    results->energy_values[5] = result.energy;
    results->force_values[5] = result.force;
    results->attractive_region_pass =
        (result.energy < 0.0) && (result.force.x < 0.0);
  }

  // Arbitrary Direction Test
  {
    Vec3<real> r = {1.0, 1.0, 1.0};
    auto result = lj.calc_force_and_energy(r);
    results->energy_values[6] = result.energy;
    results->force_values[6] = result.force;

    real r_mag = sqrt(r.squared_norm2());
    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;

    results->arbitrary_direction_pass =
        (force_dot_r < 0.0) &&
        (fabs(result.force.x - result.force.y) < tolerance) &&
        (fabs(result.force.y - result.force.z) < tolerance);
  }

  // Parameter Variation Test
  {
    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->energy_values[7] = result2.energy;
    results->force_values[7] = result2.force;

    results->parameter_variation_pass = (result1.energy != result2.energy) &&
                                        (result1.force.x != result2.force.x);
  }

  // Exact Value Check Test
  {
    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);

    results->energy_values[8] = result.energy;
    results->force_values[8] = result.force;

    real expected_energy = 4.0 * (pow(1.0 / 1.5, 12) - pow(1.0 / 1.5, 6));
    real expected_force =
        24.0 * (pow(1.0 / 1.5, 6) - 2.0 * pow(1.0 / 1.5, 12)) / 1.5;

    results->exact_value_check_pass =
        (fabs(result.energy - expected_energy) < tolerance) &&
        (fabs(result.force.x + expected_force) < tolerance) &&
        (fabs(result.force.y) < tolerance) &&
        (fabs(result.force.z) < tolerance);
  }

  // Near Cutoff Test
  {
    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);

    results->energy_values[9] = result_inside.energy;
    results->force_values[9] = result_inside.force;

    results->near_cutoff_pass =
        (result_inside.energy != 0.0) && (result_inside.force.x != 0.0) &&
        (result_outside.energy == 0.0) &&
        vec3_near(Vec3<real>{0.0, 0.0, 0.0}, result_outside.force, tolerance);
  }
}

// Helper class for CUDA error checking
class CudaErrorCheck {
public:
  static void checkAndThrow(cudaError_t err, const char *msg) {
    if (err != cudaSuccess) {
      std::string error_message =
          std::string(msg) + ": " + cudaGetErrorString(err);
      throw std::runtime_error(error_message);
    }
  }
};

// Google Test wrapper that runs the device tests
class LennardJonesCudaTest : public ::testing::Test {
protected:
  void SetUp() override {
    // Allocate device memory for results
    CudaErrorCheck::checkAndThrow(
        cudaMalloc(&d_results, sizeof(TestResults)),
        "Failed to allocate device memory for test results");
  }

  void TearDown() override {
    if (d_results) {
      cudaFree(d_results);
      d_results = nullptr;
    }
  }

  // Helper function to run tests on device and get results
  TestResults runDeviceTests() {
    TestResults h_results;

    // Clear device memory
    CudaErrorCheck::checkAndThrow(cudaMemset(d_results, 0, sizeof(TestResults)),
                                  "Failed to clear device memory");

    // Run kernel with a single thread
    lennard_jones_test_kernel<<<1, 1>>>(d_results);

    // Check for kernel launch errors
    CudaErrorCheck::checkAndThrow(cudaGetLastError(), "Kernel launch failed");

    // Wait for kernel to complete
    CudaErrorCheck::checkAndThrow(cudaDeviceSynchronize(),
                                  "Kernel execution failed");

    // Copy results back to host
    CudaErrorCheck::checkAndThrow(cudaMemcpy(&h_results, d_results,
                                             sizeof(TestResults),
                                             cudaMemcpyDeviceToHost),
                                  "Failed to copy results from device");

    return h_results;
  }

  TestResults *d_results = nullptr;
};

// Define the actual test cases
TEST_F(LennardJonesCudaTest, DeviceZeroDistance) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.zero_distance_pass)
      << "Zero distance test failed on device. Energy: "
      << results.energy_values[0] << ", Force: (" << results.force_values[0].x
      << ", " << results.force_values[0].y << ", " << results.force_values[0].z
      << ")";
}

TEST_F(LennardJonesCudaTest, DeviceBeyondCutoff) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.beyond_cutoff_pass)
      << "Beyond cutoff test failed on device. Energy: "
      << results.energy_values[1];
}

TEST_F(LennardJonesCudaTest, DeviceAtMinimum) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.at_minimum_pass)
      << "At minimum test failed on device. Energy: "
      << results.energy_values[2];
}

TEST_F(LennardJonesCudaTest, DeviceAtEquilibrium) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.at_equilibrium_pass)
      << "At equilibrium test failed on device. Energy: "
      << results.energy_values[3] << ", Force x: " << results.force_values[3].x;
}

TEST_F(LennardJonesCudaTest, DeviceRepulsiveRegion) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.repulsive_region_pass)
      << "Repulsive region test failed on device. Energy: "
      << results.energy_values[4] << ", Force x: " << results.force_values[4].x;
}

TEST_F(LennardJonesCudaTest, DeviceAttractiveRegion) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.attractive_region_pass)
      << "Attractive region test failed on device. Energy: "
      << results.energy_values[5] << ", Force x: " << results.force_values[5].x;
}

TEST_F(LennardJonesCudaTest, DeviceArbitraryDirection) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.arbitrary_direction_pass)
      << "Arbitrary direction test failed on device.";
}

TEST_F(LennardJonesCudaTest, DeviceParameterVariation) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.parameter_variation_pass)
      << "Parameter variation test failed on device.";
}

TEST_F(LennardJonesCudaTest, DeviceExactValueCheck) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.exact_value_check_pass)
      << "Exact value check test failed on device. Energy: "
      << results.energy_values[8] << ", Force x: " << results.force_values[8].x;
}

TEST_F(LennardJonesCudaTest, DeviceNearCutoff) {
  auto results = runDeviceTests();
  EXPECT_TRUE(results.near_cutoff_pass)
      << "Near cutoff test failed on device. Inside energy: "
      << results.energy_values[9];
}