#include <cmath>
#include <cuda_runtime.h>
#include <gtest/gtest.h>
#include <vector>

// Include your header files
#include "forces.cuh"
#include "pair_potentials.cuh"
#include "precision.hpp"

class CudaForceKernelTest : public ::testing::Test {
protected:
  const int BLOCK_SIZE = 1;
  const int THREADS_PER_BLOCK = 4;

  void SetUp() override {
    // Set up CUDA device
    cudaError_t err = cudaSetDevice(0);
    ASSERT_EQ(err, cudaSuccess) << "Failed to set CUDA device";
  }

  void TearDown() override {
    // Clean up any remaining GPU memory
    cudaDeviceReset();
  }

  // Helper function to check CUDA errors
  void checkCudaError(cudaError_t err, const std::string &operation) {
    ASSERT_EQ(err, cudaSuccess)
        << "CUDA error in " << operation << ": " << cudaGetErrorString(err);
  }

  // Helper function to allocate and copy data to GPU
  template <typename T>
  T *allocateAndCopyToGPU(const std::vector<T> &host_data) {
    T *device_ptr;
    size_t size = host_data.size() * sizeof(T);
    checkCudaError(cudaMalloc(&device_ptr, size), "cudaMalloc");
    checkCudaError(
        cudaMemcpy(device_ptr, host_data.data(), size, cudaMemcpyHostToDevice),
        "cudaMemcpy H2D");
    return device_ptr;
  }

  // Helper function to copy data from GPU and free GPU memory
  template <typename T>
  std::vector<T> copyFromGPUAndFree(T *device_ptr, size_t count) {
    std::vector<T> host_data(count);
    size_t size = count * sizeof(T);
    checkCudaError(
        cudaMemcpy(host_data.data(), device_ptr, size, cudaMemcpyDeviceToHost),
        "cudaMemcpy D2H");
    checkCudaError(cudaFree(device_ptr), "cudaFree");
    return host_data;
  }

  // Helper function to run the force calculation kernel
  std::pair<std::vector<real>, std::vector<real>>
  run_force_calculation(int n_particles, const std::vector<real> &positions,
                        const std::vector<real> &box_dimensions) {
    std::vector<real> forces(3 * n_particles, 0.0);
    std::vector<real> energies(n_particles, 0.0);

    real *d_positions = allocateAndCopyToGPU(positions);
    real *d_forces = allocateAndCopyToGPU(forces);
    real *d_energies = allocateAndCopyToGPU(energies);
    real *d_box_len = allocateAndCopyToGPU(box_dimensions);

    // Allocate potential on the GPU
    LennardJones h_potential(1.0, 1.0, 3.0);
    LennardJones *d_potential;
    checkCudaError(cudaMalloc(&d_potential, sizeof(LennardJones)),
                   "cudaMalloc potential");
    checkCudaError(cudaMemcpy(d_potential, &h_potential, sizeof(LennardJones),
                              cudaMemcpyHostToDevice),
                   "cudaMemcpy H2D potential");

    CAC::calc_forces_and_energies<<<BLOCK_SIZE, THREADS_PER_BLOCK>>>(
        d_positions, d_forces, d_energies, n_particles, d_box_len, d_potential);

    checkCudaError(cudaGetLastError(), "kernel launch");
    checkCudaError(cudaDeviceSynchronize(), "kernel execution");

    std::vector<real> result_forces =
        copyFromGPUAndFree(d_forces, 3 * n_particles);
    std::vector<real> result_energies =
        copyFromGPUAndFree(d_energies, n_particles);

    checkCudaError(cudaFree(d_positions), "cudaFree positions");
    checkCudaError(cudaFree(d_box_len), "cudaFree box_len");
    checkCudaError(cudaFree(d_potential), "cudaFree potential");

    return {result_forces, result_energies};
  }
};

TEST_F(CudaForceKernelTest, BasicFunctionalityTest) {
  const int n_particles = 2;
  const real tolerance = 1e-5;

  // Set up test data - simple 2x2 grid of particles
  std::vector<real> positions = {
      0.0, 0.0, 0.0, // particle 0
      1.0, 0.0, 0.0, // particle 1
  };

  std::vector<real> box_dimensions = {10.0, 10.0, 10.0};

  auto [result_forces, result_energies] =
      run_force_calculation(n_particles, positions, box_dimensions);

  // Verify results - forces should be non-zero and energies should be
  // calculated
  bool has_nonzero_force = false;
  bool has_nonzero_energy = false;

  for (int i = 0; i < 3 * n_particles; i++) {
    if (std::abs(result_forces[i]) > tolerance) {
      has_nonzero_force = true;
      break;
    }
  }

  for (int i = 0; i < n_particles; i++) {
    if (std::abs(result_energies[i]) > tolerance) {
      has_nonzero_energy = true;
      break;
    }
  }

  EXPECT_TRUE(has_nonzero_force)
      << "Expected non-zero forces between particles";
  EXPECT_TRUE(has_nonzero_energy)
      << "Expected non-zero energies for particles ";
}
//
// TEST_F(CudaKernelTest, PeriodicBoundaryConditionsTest) {
//   const int n_particles = 2;
//   const real tolerance = 1e-5;
//
//   // Place particles near opposite edges of a small box
//   std::vector<real> positions = {
//       0.1, 0.0, 0.0, // particle 0 near left edge
//       4.9, 0.0, 0.0  // particle 1 near right edge
//   };
//   std::vector<real> box_dimensions = {5.0, 5.0, 5.0}; // Small box to test
//   PBC
//
//   auto [result_forces, result_energies] =
//       run_force_calculation(n_particles, &positions, &box_dimensions);
//
//   // With PBC, particles should interact as if they're close (distance ~0.2)
//   // rather than far apart (distance ~4.8)
//   EXPECT_GT(std::abs(result_forces[0]), tolerance)
//       << "Expected significant force due to PBC";
//   EXPECT_GT(std::abs(result_energies[0]), tolerance)
//       << "Expected significant energy due to PBC";
// }

// TEST_F(CudaForceKernelTest, SingleParticleTest) {
//   const int n_particles = 1;
//
//   std::vector<real> positions = {0.0, 0.0, 0.0};
//   std::vector<real> box_dimensions = {10.0, 10.0, 10.0};
//
//   auto [result_forces, result_energies] =
//       run_force_calculation(n_particles, positions, box_dimensions);
//   // Single particle should have zero force and energy
//   EXPECT_NEAR(result_forces[0], 0.0, 1e-10);
//   EXPECT_NEAR(result_forces[1], 0.0, 1e-10);
//   EXPECT_NEAR(result_forces[2], 0.0, 1e-10);
//   EXPECT_NEAR(result_energies[0], 0.0, 1e-10);
// }

// TEST_F(CudaKernelTest, ForceSymmetryTest) {
//   const int n_particles = 2;
//   const real tolerance = 1e-5;
//
//   std::vector<real> positions = {
//       0.0, 0.0, 0.0, // particle 0
//       1.5, 0.0, 0.0  // particle 1
//   };
//   std::vector<real> box_dimensions = {10.0, 10.0, 10.0};
//
//   auto [result_forces, result_energies] =
//       run_force_calculation(n_particles, &positions, &box_dimensions);
//
//   // Newton's third law: forces should be equal and opposite
//   EXPECT_NEAR(result_forces[0], -result_forces[3], tolerance)
//       << "Force x-components should be opposite";
//   EXPECT_NEAR(result_forces[1], -result_forces[4], tolerance)
//       << "Force y-components should be opposite";
//   EXPECT_NEAR(result_forces[2], -result_forces[5], tolerance)
//       << "Force z-components should be opposite";
//
//   // Energies should be equal for symmetric particles
//   EXPECT_NEAR(result_energies[0], result_energies[1], tolerance)
//       << "Energies should be equal";
// }