eightyseven 0.2.0-alpha

use arraystring::{typenum::U5, ArrayString};

type Vec3 = [f32; 3];

/// 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 {
        let n = self.natoms() as f32;
        self.atoms
            .iter()
            .map(|atom| atom.position)
            .fold(Default::default(), |acc: Vec3, position| {
                [
                    acc[0] + position[0],
                    acc[1] + position[1],
                    acc[2] + position[2],
                ]
            })
            .map(|v| v / n)
    }

    /// 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[0] -= center[0];
            atom.position[1] -= center[1];
            atom.position[2] -= center[2];
        });
    }

    /// 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 `[f32; 3]`s in `v1 v2 v3` order.
    #[must_use]
    pub const fn as_vecs(&self) -> [Vec3; 3] {
        match *self {
            BoxVecs::Short([v1x, v2y, v3z]) => [
                [v1x, 0.0, 0.0], // v1
                [0.0, v2y, 0.0], // v2
                [0.0, 0.0, v3z], // v3
            ],
            BoxVecs::Full([v1x, v2y, v3z, v1y, v1z, v2x, v2z, v3x, v3y]) => [
                [v1x, v1y, v1z], // v1
                [v2x, v2y, v2z], // v2
                [v3x, v3y, v3z], // v3
            ],
        }
    }
}

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 = glam::Vec3::new(3.9875, 3.9760, 2.7035);
        assert!(glam::Vec3::from_array(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!(glam::Vec3::from_array(structure.center()).abs_diff_eq(glam::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(),
            [[10.0, 0.0, 0.0], [0.0, 20.0, 0.0], [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(),
            [[10.0, 40.0, 50.0], [60.0, 20.0, 70.0], [80.0, 90.0, 30.0]]
        );
    }
}