eightyseven 0.1.5

use arraystring::{typenum::U5, ArrayString};
pub use glam::Vec3;

/// The contents of a `gro` file.
#[derive(Debug, Clone, PartialEq)]
pub struct Structure {
    pub title: String,
    pub atoms: Vec<Atom>,
    pub boxvecs: BoxVecs,
}

impl Structure {
    /// Return the center of the positions of all atoms in the structure.
    ///
    /// All atoms are weighed equally, so this is not necessarily the center of mass. In other
    /// words, the average position of all the atoms is returned.
    ///
    /// If there are not atoms in the system, a zero vector is returned.
    #[must_use]
    pub fn center(&self) -> Vec3 {
        self.atoms.iter().map(|atom| atom.position).sum::<Vec3>() / self.natoms() as f32
    }

    /// Centers all atoms of this [`Structure`] such that the [`Structure::center`] is zero.
    pub fn translate_to_center(&mut self) {
        let center = self.center();
        self.atoms
            .iter_mut()
            .for_each(|atom| atom.position -= center);
    }

    /// Returns the number of atoms in this [`Structure`].
    #[must_use]
    pub fn natoms(&self) -> usize {
        self.atoms.len()
    }
}

pub type ResNum = u32;
pub type AtomNum = u32;
pub type ResName = ArrayString<U5>;
pub type AtomName = ArrayString<U5>;

#[derive(Debug, Default, Clone, Copy, PartialEq)]
pub struct Atom {
    // 5 positions, integer.
    /// Residue number.
    pub resnum: ResNum,
    // 5 characters.
    /// Residue name.
    pub resname: ResName,
    // 5 characters.
    /// Atom name.
    pub atomname: AtomName,
    // 5 positions, integer.
    /// Atom number.
    pub atomnum: AtomNum,
    // In nm, x y z in 3 columns, each 8 positions with 3 decimal places.
    /// Position (nm).
    pub position: Vec3,
    // In nm/ps, x y z in 3 columns, each 8 positions with 4 decimal places.
    /// Velocity (nm/ps or km/s).
    pub velocity: Vec3,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum BoxVecs {
    /// Short 3-value box vector (`v1(x) v2(y) v3(z)`).
    Short([f32; 3]),
    /// Full 9-value box vectors (`v1(x) v2(y) v3(z) v1(y) v1(z) v2(x) v2(z) v3(x) v3(y)`).
    Full([f32; 9]),
}

impl BoxVecs {
    /// Return these [`BoxVecs`] as an array in `gro` order.
    ///
    /// `v1(x) v2(y) v3(z) v1(y) v1(z) v2(x) v2(z) v3(x) v3(y)`
    ///
    /// If only the first three positions are known, the last six positions will be set to zero.
    #[must_use]
    pub const fn as_array(&self) -> [f32; 9] {
        match *self {
            BoxVecs::Short([v1x, v2y, v3z]) => [v1x, v2y, v3z, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            BoxVecs::Full(bv) => bv,
        }
    }

    /// Return these [`BoxVecs`] as three [`Vec3`]s in `v1 v2 v3` order.
    #[must_use]
    pub const fn as_vecs(&self) -> [Vec3; 3] {
        match *self {
            BoxVecs::Short([v1x, v2y, v3z]) => [
                Vec3::new(v1x, 0.0, 0.0),
                Vec3::new(0.0, v2y, 0.0),
                Vec3::new(0.0, 0.0, v3z),
            ],
            BoxVecs::Full([v1x, v2y, v3z, v1y, v1z, v2x, v2z, v3x, v3y]) => [
                Vec3::new(v1x, v1y, v1z),
                Vec3::new(v2x, v2y, v2z),
                Vec3::new(v3x, v3y, v3z),
            ],
        }
    }
}

impl std::fmt::Display for BoxVecs {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        fn list<const N: usize>(vs: &[f32; N]) -> String {
            vs.map(|v| v.to_string()).join(" ")
        }

        f.write_str(&match self {
            BoxVecs::Short(vs) => list(vs),
            BoxVecs::Full(vs) => list(vs),
        })
    }
}

#[cfg(test)]
mod tests {
    use std::io;

    use super::*;
    use crate::reader::ReadGro;

    const EPS: f32 = 0.0001; // For approximate float comparisons.

    #[test]
    fn center() -> io::Result<()> {
        let structure = Structure::open_gro(crate::tests::PATH)?;
        let center = Vec3::new(3.9875, 3.9760, 2.7035);
        assert!(structure.center().abs_diff_eq(center, EPS));
        Ok(())
    }

    #[test]
    fn translate_to_center() -> io::Result<()> {
        let mut structure = Structure::open_gro(crate::tests::PATH)?;
        structure.translate_to_center();
        assert!(structure.center().abs_diff_eq(Vec3::ZERO, EPS));
        Ok(())
    }

    #[test]
    fn boxvecs_short() {
        let short = BoxVecs::Short([10.0, 20.0, 30.0]);

        assert_eq!(
            short.as_array(),
            [10.0, 20.0, 30.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        );
        assert_eq!(
            short.as_vecs(),
            [
                Vec3::new(10.0, 0.0, 0.0),
                Vec3::new(0.0, 20.0, 0.0),
                Vec3::new(0.0, 0.0, 30.0)
            ]
        );
    }

    #[test]
    fn boxvecs_long() {
        let full = BoxVecs::Full([10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0]);

        assert_eq!(
            full.as_array(),
            [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0]
        );
        assert_eq!(
            full.as_vecs(),
            [
                Vec3::new(10.0, 40.0, 50.0),
                Vec3::new(60.0, 20.0, 70.0),
                Vec3::new(80.0, 90.0, 30.0)
            ]
        );
    }
}