#ifndef GAIA_CONSTANT_PH_CPHMD_HPP
#define GAIA_CONSTANT_PH_CPHMD_HPP

#include <vector>
#include <cmath>
#include <random>
#include <algorithm>
#include <string>
#include <iostream>

namespace gaia {
namespace constant_ph {

struct Vec3 {
    double x, y, z;
    Vec3(double x=0, double y=0, double z=0) : x(x), y(y), z(z) {}
    Vec3 operator+(const Vec3& o) const { return Vec3(x+o.x, y+o.y, z+o.z); }
    Vec3 operator-(const Vec3& o) const { return Vec3(x-o.x, y-o.y, z-o.z); }
    Vec3 operator*(double s) const { return Vec3(x*s, y*s, z*s); }
    double dot(const Vec3& o) const { return x*o.x + y*o.y + z*o.z; }
    double norm2() const { return x*x + y*y + z*z; }
    double norm() const { return sqrt(norm2()); }
};

struct Residue {
    int id;
    std::string name;
    int protonation_state;
    double pKa_ref;
    double pKa_shift;
    Vec3 position;
    double lambda;
    double lambda_vel;
    double lambda_mass;
};

class ConstantpH {
public:
    ConstantpH(double pH=7.0, double temp=300.0) : pH(pH), kT(0.001987204258*temp) {}
    
    double compute_pKa_shift(const Residue& res, const std::vector<Vec3>& charges,
                             const std::vector<Vec3>& positions) {
        double E_shift = 0.0;
        for (size_t i = 0; i < charges.size(); i++) {
            Vec3 dr = res.position - positions[i];
            double r = dr.norm();
            if (r < 0.1) continue;
            E_shift += 332.0 * charges[i].x / (4.0 * r);
        }
        return E_shift / (2.303 * kT);
    }
    
    bool attempt_titration(Residue& res, const std::vector<Vec3>& charges,
                           const std::vector<Vec3>& positions) {
        res.pKa_shift = compute_pKa_shift(res, charges, positions);
        double pKa_eff = res.pKa_ref + res.pKa_shift;
        double delta_G = 2.303 * kT * (pH - pKa_eff);
        if (delta_G < 0 || exp(-delta_G/kT) > uniform_random()) {
            res.protonation_state = 1 - res.protonation_state;
            return true;
        }
        return false;
    }
    
    void run_titration(std::vector<Residue>& residues,
                       const std::vector<Vec3>& charges,
                       const std::vector<Vec3>& positions,
                       int steps, int freq=1000) {
        for (int step = 0; step < steps; step++) {
            if (step % freq == 0 && step > 0) {
                for (auto& res : residues) {
                    if (attempt_titration(res, charges, positions)) {
                        std::cout << "Residue " << res.id << " titrated at step " << step << "\n";
                    }
                }
            }
        }
    }

private:
    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);
    }
    double pH, kT;
};

} // namespace constant_ph
} // namespace gaia

#endif