// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// ForceField.hpp — 分子力場計算 (CHEM-1)
// Lennard-Jones, Coulomb, 結合伸縮ポテンシャルと最急降下法による構造最適化

#ifndef CALX_CHEM_FORCEFIELD_HPP
#define CALX_CHEM_FORCEFIELD_HPP

#include <vector>
#include <array>
#include <cmath>
#include <cstddef>
#include <stdexcept>
#include <algorithm>

namespace calx::chem {

    // 原子
    template<typename T>
    struct Atom {
        T x, y, z;        // 座標
        T charge;          // 部分電荷
        T sigma, epsilon;  // Lennard-Jones パラメータ
    };

    // 結合
    template<typename T>
    struct Bond {
        size_t atom1, atom2;
        T k;   // バネ定数
        T r0;  // 平衡結合長
    };

    // 分子力場クラス
    template<typename T>
    class ForceField {
        std::vector<Atom<T>> atoms_;
        std::vector<Bond<T>> bonds_;

        // Coulomb 定数 (kcal/mol·Å·e^-2, AMBER 単位系)
        static constexpr double COULOMB_CONST = 332.0636;

    public:
        void addAtom(const Atom<T>& a) { atoms_.push_back(a); }

        void addBond(const Bond<T>& b) {
            if (b.atom1 >= atoms_.size() || b.atom2 >= atoms_.size())
                throw std::out_of_range("addBond: atom index out of range");
            bonds_.push_back(b);
        }

        const std::vector<Atom<T>>& atoms() const { return atoms_; }
        const std::vector<Bond<T>>& bonds() const { return bonds_; }

        // -----------------------------------------------------------
        // エネルギー計算
        // -----------------------------------------------------------

        // Lennard-Jones ポテンシャル (非結合対)
        // V_LJ = 4ε [(σ/r)^12 - (σ/r)^6]
        // Lorentz-Berthelot 混合則: σ_ij = (σ_i + σ_j)/2, ε_ij = sqrt(ε_i * ε_j)
        T lennardJonesEnergy() const {
            T energy = T(0);
            const size_t n = atoms_.size();

            // 結合原子ペアをスキップするためのセット
            auto isBonded = [this](size_t i, size_t j) -> bool {
                for (const auto& b : bonds_) {
                    if ((b.atom1 == i && b.atom2 == j) ||
                        (b.atom1 == j && b.atom2 == i))
                        return true;
                }
                return false;
            };

            for (size_t i = 0; i < n; ++i) {
                for (size_t j = i + 1; j < n; ++j) {
                    if (isBonded(i, j)) continue;

                    T dx = atoms_[j].x - atoms_[i].x;
                    T dy = atoms_[j].y - atoms_[i].y;
                    T dz = atoms_[j].z - atoms_[i].z;
                    T r2 = dx * dx + dy * dy + dz * dz;

                    if (r2 < T(1e-12)) continue; // 重なり回避

                    T sig = (atoms_[i].sigma + atoms_[j].sigma) / T(2);
                    T eps = std::sqrt(atoms_[i].epsilon * atoms_[j].epsilon);

                    T s2r2 = (sig * sig) / r2;
                    T s6 = s2r2 * s2r2 * s2r2;
                    energy += T(4) * eps * (s6 * s6 - s6);
                }
            }
            return energy;
        }

        // Coulomb ポテンシャル (非結合対)
        // V_C = k_e * q_i * q_j / r
        T coulombEnergy() const {
            T energy = T(0);
            const size_t n = atoms_.size();

            auto isBonded = [this](size_t i, size_t j) -> bool {
                for (const auto& b : bonds_) {
                    if ((b.atom1 == i && b.atom2 == j) ||
                        (b.atom1 == j && b.atom2 == i))
                        return true;
                }
                return false;
            };

            for (size_t i = 0; i < n; ++i) {
                for (size_t j = i + 1; j < n; ++j) {
                    if (isBonded(i, j)) continue;

                    T dx = atoms_[j].x - atoms_[i].x;
                    T dy = atoms_[j].y - atoms_[i].y;
                    T dz = atoms_[j].z - atoms_[i].z;
                    T r = std::sqrt(dx * dx + dy * dy + dz * dz);

                    if (r < T(1e-12)) continue;

                    energy += static_cast<T>(COULOMB_CONST)
                            * atoms_[i].charge * atoms_[j].charge / r;
                }
            }
            return energy;
        }

        // 結合伸縮ポテンシャル (調和振動子)
        // V_bond = (1/2) k (r - r0)^2
        T bondEnergy() const {
            T energy = T(0);
            for (const auto& b : bonds_) {
                T dx = atoms_[b.atom2].x - atoms_[b.atom1].x;
                T dy = atoms_[b.atom2].y - atoms_[b.atom1].y;
                T dz = atoms_[b.atom2].z - atoms_[b.atom1].z;
                T r = std::sqrt(dx * dx + dy * dy + dz * dz);
                T dr = r - b.r0;
                energy += T(0.5) * b.k * dr * dr;
            }
            return energy;
        }

        // 全エネルギー
        T totalEnergy() const {
            return lennardJonesEnergy() + coulombEnergy() + bondEnergy();
        }

        // -----------------------------------------------------------
        // 力の計算 (エネルギーの負の勾配)
        // -----------------------------------------------------------
        std::vector<std::array<T, 3>> forces() const {
            const size_t n = atoms_.size();
            std::vector<std::array<T, 3>> f(n, {T(0), T(0), T(0)});

            auto isBonded = [this](size_t i, size_t j) -> bool {
                for (const auto& b : bonds_) {
                    if ((b.atom1 == i && b.atom2 == j) ||
                        (b.atom1 == j && b.atom2 == i))
                        return true;
                }
                return false;
            };

            // Lennard-Jones の力
            for (size_t i = 0; i < n; ++i) {
                for (size_t j = i + 1; j < n; ++j) {
                    if (isBonded(i, j)) continue;

                    T dx = atoms_[j].x - atoms_[i].x;
                    T dy = atoms_[j].y - atoms_[i].y;
                    T dz = atoms_[j].z - atoms_[i].z;
                    T r2 = dx * dx + dy * dy + dz * dz;

                    if (r2 < T(1e-12)) continue;

                    T sig = (atoms_[i].sigma + atoms_[j].sigma) / T(2);
                    T eps = std::sqrt(atoms_[i].epsilon * atoms_[j].epsilon);

                    T s2r2 = (sig * sig) / r2;
                    T s6 = s2r2 * s2r2 * s2r2;

                    // F_LJ = 24ε/r² [2(σ/r)^12 - (σ/r)^6] * r_vec
                    T fMag = T(24) * eps / r2 * (T(2) * s6 * s6 - s6);

                    f[i][0] -= fMag * dx; f[i][1] -= fMag * dy; f[i][2] -= fMag * dz;
                    f[j][0] += fMag * dx; f[j][1] += fMag * dy; f[j][2] += fMag * dz;
                }
            }

            // Coulomb の力
            for (size_t i = 0; i < n; ++i) {
                for (size_t j = i + 1; j < n; ++j) {
                    if (isBonded(i, j)) continue;

                    T dx = atoms_[j].x - atoms_[i].x;
                    T dy = atoms_[j].y - atoms_[i].y;
                    T dz = atoms_[j].z - atoms_[i].z;
                    T r2 = dx * dx + dy * dy + dz * dz;
                    T r = std::sqrt(r2);

                    if (r < T(1e-12)) continue;

                    // F_C = -k_e * q_i * q_j / r² * r_hat
                    //      = -k_e * q_i * q_j / r³ * r_vec
                    T fMag = -static_cast<T>(COULOMB_CONST)
                           * atoms_[i].charge * atoms_[j].charge / (r2 * r);

                    f[i][0] -= fMag * dx; f[i][1] -= fMag * dy; f[i][2] -= fMag * dz;
                    f[j][0] += fMag * dx; f[j][1] += fMag * dy; f[j][2] += fMag * dz;
                }
            }

            // 結合の力
            for (const auto& b : bonds_) {
                T dx = atoms_[b.atom2].x - atoms_[b.atom1].x;
                T dy = atoms_[b.atom2].y - atoms_[b.atom1].y;
                T dz = atoms_[b.atom2].z - atoms_[b.atom1].z;
                T r = std::sqrt(dx * dx + dy * dy + dz * dz);

                if (r < T(1e-12)) continue;

                // F_bond = -k(r - r0) * r_hat
                T fMag = -b.k * (r - b.r0) / r;

                f[b.atom1][0] -= fMag * dx; f[b.atom1][1] -= fMag * dy; f[b.atom1][2] -= fMag * dz;
                f[b.atom2][0] += fMag * dx; f[b.atom2][1] += fMag * dy; f[b.atom2][2] += fMag * dz;
            }

            return f;
        }

        // -----------------------------------------------------------
        // エネルギー最小化 (最急降下法)
        // -----------------------------------------------------------
        void minimize(size_t maxIter = 1000, T stepSize = T(0.001)) {
            for (size_t iter = 0; iter < maxIter; ++iter) {
                auto f = forces();

                // 最大力の大きさを計算 (収束判定用)
                T maxForce = T(0);
                for (size_t i = 0; i < atoms_.size(); ++i) {
                    T fMag2 = f[i][0] * f[i][0] + f[i][1] * f[i][1] + f[i][2] * f[i][2];
                    maxForce = std::max(maxForce, fMag2);
                }

                // 収束判定: 最大力が十分小さい
                if (maxForce < T(1e-10)) break;

                // 座標更新: x += stepSize * force
                for (size_t i = 0; i < atoms_.size(); ++i) {
                    atoms_[i].x += stepSize * f[i][0];
                    atoms_[i].y += stepSize * f[i][1];
                    atoms_[i].z += stepSize * f[i][2];
                }
            }
        }
    };

} // namespace calx::chem

#endif // CALX_CHEM_FORCEFIELD_HPP
