#ifndef GAIA_AMOEBA_CG_SOLVER_HPP
#define GAIA_AMOEBA_CG_SOLVER_HPP

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

namespace gaia {
namespace amoeba {

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; }
};

struct Atom {
    Vec3 pos;
    Vec3 mu_prev;
    double alpha;
    int element;
};

class ConjugateGradientSolver {
public:
    ConjugateGradientSolver(double tolerance = 1e-6, int max_iter = 20)
        : tolerance(tolerance), max_iter(max_iter) {}

    std::vector<Vec3> solve(const std::vector<Atom>& atoms,
                            const std::vector<Vec3>& E_field) {
        int n = (int)atoms.size();
        std::vector<Vec3> mu(n);
        std::vector<Vec3> r(n);
        std::vector<Vec3> p(n);
        
        for (int i = 0; i < n; i++) {
            mu[i] = atoms[i].mu_prev;
            r[i] = E_field[i] - compute_T_mu(atoms, mu, i);
            p[i] = r[i];
        }
        
        double rho = dot_product(r, r);
        
        for (int iter = 0; iter < max_iter; iter++) {
            std::vector<Vec3> Ap(n);
            for (int i = 0; i < n; i++) {
                Ap[i] = compute_T_p(atoms, p, i);
            }
            
            double pAp = dot_product(p, Ap);
            if (pAp < 1e-12) break;
            
            double alpha = rho / pAp;
            
            for (int i = 0; i < n; i++) {
                mu[i] = mu[i] + p[i] * alpha;
            }
            
            for (int i = 0; i < n; i++) {
                r[i] = r[i] - Ap[i] * alpha;
            }
            
            double rho_new = dot_product(r, r);
            double beta = (rho_new - rho) / rho;
            
            for (int i = 0; i < n; i++) {
                p[i] = r[i] + p[i] * beta;
            }
            
            rho = rho_new;
            if (std::sqrt(rho) < tolerance) break;
        }
        
        return mu;
    }

private:
    Vec3 compute_T_mu(const std::vector<Atom>& atoms,
                      const std::vector<Vec3>& mu, int i) {
        Vec3 result;
        const double damping = 0.39;
        
        for (int j = 0; j < (int)atoms.size(); j++) {
            if (i == j) continue;
            Vec3 dr = atoms[i].pos - atoms[j].pos;
            double r2 = dr.norm2();
            double r = std::sqrt(r2);
            if (r < 1e-12) continue;
            
            double alpha_i = atoms[i].alpha;
            double alpha_j = atoms[j].alpha;
            double u = r / std::pow(alpha_i * alpha_j, 1.0/6.0);
            double damp = 1.0 - std::exp(-damping * u * u * u);
            double factor = damp * alpha_j / (r2 * r);
            result = result + dr * (mu[j].dot(dr) * 3.0 * factor) - mu[j] * factor;
        }
        return result;
    }
    
    Vec3 compute_T_p(const std::vector<Atom>& atoms,
                     const std::vector<Vec3>& p, int i) {
        Vec3 result;
        const double damping = 0.39;
        
        for (int j = 0; j < (int)atoms.size(); j++) {
            if (i == j) continue;
            Vec3 dr = atoms[i].pos - atoms[j].pos;
            double r2 = dr.norm2();
            double r = std::sqrt(r2);
            if (r < 1e-12) continue;
            
            double alpha_i = atoms[i].alpha;
            double alpha_j = atoms[j].alpha;
            double u = r / std::pow(alpha_i * alpha_j, 1.0/6.0);
            double damp = 1.0 - std::exp(-damping * u * u * u);
            double factor = damp * alpha_j / (r2 * r);
            result = result + dr * (p[j].dot(dr) * 3.0 * factor) - p[j] * factor;
        }
        return result;
    }
    
    double dot_product(const std::vector<Vec3>& a, const std::vector<Vec3>& b) {
        double sum = 0.0;
        for (size_t i = 0; i < a.size(); i++) {
            sum += a[i].dot(b[i]);
        }
        return sum;
    }
    
    double tolerance;
    int max_iter;
};

} // namespace amoeba
} // namespace gaia

#endif