#ifndef GAIA_GCMC_NCMC_HPP
#define GAIA_GCMC_NCMC_HPP

#include <vector>
#include <cmath>
#include <random>

namespace gaia {
namespace gcmc {

struct Vec3 {
    double x, y, z;
    Vec3(double x=0, double y=0, double z=0) : x(x), y(y), z(z) {}
};

struct Water {
    Vec3 O, H1, H2;
    double lambda;
    Water(Vec3 pos, double l=0.0) : lambda(l) {
        double bond = 0.9572;
        double angle = 104.52 * 3.14159 / 180.0;
        O = pos;
        H1 = Vec3(pos.x + bond * cos(angle/2), pos.y + bond * sin(angle/2), pos.z);
        H2 = Vec3(pos.x + bond * cos(-angle/2), pos.y + bond * sin(-angle/2), pos.z);
    }
    void set_lambda(double l) { lambda = l; }
};

class NCMCWater {
public:
    NCMCWater(int n_steps=20, double overlap=2.0) : n_steps(n_steps), overlap(overlap) {}
    
    bool insert(const std::vector<Vec3>& protein, const std::vector<Water>& waters,
                const Vec3& pos, double kT) {
        const double mu = -6.09;
        for (const auto& w : waters) {
            double dx = pos.x - w.O.x, dy = pos.y - w.O.y, dz = pos.z - w.O.z;
            if (dx*dx + dy*dy + dz*dz < overlap*overlap) return false;
        }
        for (const auto& atom : protein) {
            double dx = pos.x - atom.x, dy = pos.y - atom.y, dz = pos.z - atom.z;
            if (dx*dx + dy*dy + dz*dz < 1.0) return false;
        }
        Water water(pos, 0.0);
        double U_init = compute_potential(protein, waters, water);
        double work = 0.0;
        for (int s = 0; s < n_steps; s++) {
            double lambda = (s + 1.0) / n_steps;
            water.set_lambda(lambda);
            double U_cur = compute_potential(protein, waters, water);
            work += U_cur - U_init;
            U_init = U_cur;
        }
        double delta = work - mu;
        if (delta < 0 || exp(-delta / kT) > uniform_random()) return true;
        return false;
    }

private:
    double compute_potential(const std::vector<Vec3>& protein,
                             const std::vector<Water>& waters,
                             const Water& water) {
        double U = 0.0;
        for (const auto& atom : protein) {
            double dx = water.O.x - atom.x, dy = water.O.y - atom.y, dz = water.O.z - atom.z;
            double r2 = dx*dx + dy*dy + dz*dz;
            if (r2 < 0.01) continue;
            double r6 = r2 * r2 * r2;
            U += 4.0 * 0.1 * (1.0/(r6*r6) - 1.0/r6) * water.lambda;
        }
        for (const auto& atom : protein) {
            double dx = water.O.x - atom.x, dy = water.O.y - atom.y, dz = water.O.z - atom.z;
            double r = sqrt(dx*dx + dy*dy + dz*dz);
            if (r < 0.01) continue;
            U += 332.0 * (-0.82) / r * water.lambda;
        }
        return U;
    }
    double uniform_random() {
        static std::random_device rd;
        static std::mt19937 gen(rd());
        static std::uniform_real_distribution<double> dist(0.0, 1.0);
        return dist(gen);
    }
    int n_steps;
    double overlap;
};

} // namespace gcmc
} // namespace gaia

#endif