// [[file:../elastic.note::55426a56][55426a56]]
//! # elastic energy and force
//!
//! For computing the energy and force of a center node for elastic
//! connections with its immediate neighbors.
//!
//! # Example
//!
//! ```ignore
//! // build connections
//! let mut connections = vec![];
//! let pi: Vector3f = positions[i].into();
//! for j in 0..n {
//!     let pj: Vector3f = positions[j].into();
//!     let displace = pj - pi;
//!     let rij = displace.norm();
//!     let spring = Spring::new(rij.powi(-2)).with_length(rij);
//!     let connect = Connection::new(spring, displace);
//!     connections.push(connect);
//! }
//! 
//! // compute energy and force for center node i
//! let (e, f) = compute_elastic_contribution(connections);
//! ```
#![deny(warnings)]

use super::*;
// 55426a56 ends here

// [[file:../elastic.note::4d6619eb][4d6619eb]]
impl Spring {
    /// Construct an `Spring` for elastic interaction between node `i` and node
    /// `j`
    ///
    /// # Parameters
    ///
    /// * wij: the weight (elastic constant) of spring
    pub fn new(wij: f64) -> Self {
        assert!(wij.is_sign_positive(), "invalid {wij}");
        Self { dij: 0.0, wij }
    }

    /// Set optional natural length of spring. The default value is 0.
    pub fn with_length(mut self, dij: f64) -> Self {
        assert!(dij.is_sign_positive(), "invalid {dij}");
        self.dij = dij;
        self
    }

    /// Returns energy and the force contribution of its neighbor j to center i.
    ///
    /// `displace` is the displacement vector from center i to its neighbor j.
    /// `distance` is optional, if not set, it will be evaluated from
    /// displacement vector.
    pub fn energy_force(&self, displace: impl Into<Vector3f>, distance: impl Into<Option<f64>>) -> (f64, Vector3f) {
        let pij = displace.into();
        let rij = distance.into().unwrap_or_else(|| pij.norm());
        let wij = self.wij;
        let dij = self.dij;
        let eij = 0.5 * wij * (rij - dij).powi(2);
        let fij = wij * (rij - dij) * pij / rij;
        (eij, fij)
    }
}

#[test]
fn test_spring() {
    let spring = Spring::new(1.0);
    let (e, f) = spring.energy_force([1.0; 3], None);
    approx::assert_relative_eq!(e, 1.5, epsilon = 1E-4);
    let f_expected: Vector3f = [1.0; 3].into();
    approx::assert_relative_eq!(f, f_expected, epsilon = 1E-4);

    let spring = Spring::new(1.0).with_length(3f64.sqrt());
    let (_, f) = spring.energy_force([1.0; 3], None);
    assert_eq!(f.norm(), 0.0);
}
// 4d6619eb ends here

// [[file:../elastic.note::6df983e7][6df983e7]]
impl Connection {
    /// Create an elastic `Connection` for central atom with its neighbor
    ///
    /// # Parameters
    ///
    /// * spring: elastic Spring
    /// * displace: the displacement vector from center to its neighbor
    pub fn new(spring: Spring, displace: impl Into<Array3>) -> Self {
        Self {
            spring,
            displace: displace.into(),
            distance: None,
        }
    }

    /// Set displacement vector from the center to one of its
    /// immedidate neighbors.
    pub fn set_displace(&mut self, displace: impl Into<Array3>) {
        self.displace = displace.into();
    }

    /// Set distance from the center to one of its immediate
    /// neighbors.
    pub fn set_distance(&mut self, distance: f64) {
        self.distance = distance.into();
    }
}
// 6df983e7 ends here

// [[file:../elastic.note::fd1cb865][fd1cb865]]
/// Compute elastic energy and force contribution of a center node from its
/// elastic connections with its neighbors
pub fn compute_elastic_contribution<'a, I>(env: I) -> (f64, Vector3f)
where
    I: IntoIterator<Item = &'a Connection>,
{
    let (esum, fsum) = env
        .into_iter()
        .map(|connection| connection.spring.energy_force(connection.displace, connection.distance))
        .fold((0.0, [0.0; 3].into()), |acc, x| (acc.0 + x.0, acc.1 + x.1));

    // each pair contributes a half of the spring energy
    (esum / 2.0, fsum)
}


/// Compute elastic energy/force between image 1 in `position1` and
/// image 2 in `positions2`
pub fn compute_elastic_energy_and_force_distance_space(
    positions1: &[Array3],
    positions2: &[Array3],
) -> (f64, Vec<Array3>) {
    let n = positions1.len();
    assert_eq!(positions2.len(), n);

    let mut energy = 0.0;
    let mut forces: Vec<Array3> = vec![];
    let mut connections_list = create_pairwise_connections(&positions1);

    for i in 0..n {
        for (k, j) in (0..i).chain(i + 1..n).enumerate() {
            let pi: Vector3f = positions2[i].into();
            let pj: Vector3f = positions2[j].into();
            let displace = pj - pi;
            let rij = displace.norm();
            let connection = &mut connections_list[i][k];
            connection.set_distance(rij);
            connection.set_displace(displace);
        }
        let connections = &connections_list[i];
        let (e, f) = compute_elastic_contribution(connections);
        energy += e;
        forces.push(f.into());
    }
    (energy, forces)
}
// fd1cb865 ends here

// [[file:../elastic.note::c78b9921][c78b9921]]
fn create_pairwise_connections(positions: &[Array3]) -> Vec<Vec<Connection>> {
    let mut connections_list = vec![];
    let n = positions.len();
    for i in 0..n {
        // ignore when i == j
        let mut connections = vec![];
        for j in (0..i).chain(i + 1..n) {
            let pi: Vector3f = positions[i].into();
            let pj: Vector3f = positions[j].into();
            let displace = pj - pi;
            let rij = displace.norm();
            let spring = Spring::new(rij.powi(-2)).with_length(rij);
            let connect = Connection::new(spring, displace);
            connections.push(connect);
        }
        connections_list.push(connections);
    }
    connections_list
}

#[test]
fn test_distance_potential() -> Result<()> {
    let mols = gchemol::io::read("./tests/files/md-traj-small.xyz")?.collect_vec();
    let mol = &mols[0];
    let positions1 = mol.positions().collect_vec();
    // compute elastic interaction with mols[0]
    let mol = &mols[1];
    let positions2 = mol.positions().collect_vec();
    let (energy, forces) = compute_elastic_energy_and_force_distance_space(&positions1, &positions2);

    #[rustfmt::skip]
    let f_expected = [[ -0.002955,  -0.006938,   0.019315],
                      [ -0.000554,  -0.005471,  -0.001267],
                      [ -0.007219,  -0.005921,  -0.018780],
                      [  0.015299,   0.022290,   0.002067],
                      [ -0.004571,  -0.003960,  -0.001336]];
    approx::assert_relative_eq!(forces.to_matrix(), f_expected.to_matrix(), epsilon = 1E-4);
    approx::assert_relative_eq!(energy, 0.001980439069872339, epsilon = 1E-4);

    Ok(())
}
// c78b9921 ends here