eightyseven 0.1.5

use std::io;
use std::path::Path;

use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

use crate::structure::{Atom, Structure};
use crate::GRO_INTEGER_LIMIT;

/// A trait for writing out a value as `gro` using parallel formatting.
///
/// For multi-threaded formatting, see [`WriteGroPar`].
///
/// `WriteGro` implements robust default implementations for writing to generic writers
/// ([`WriteGro::write`]) and to files ([`WriteGro::save_gro`]).
///
/// To implement this trait for a type, [`WriteGro::title`], [`WriteGro::natoms`],
/// [`WriteGro::atoms`], [`WriteGro::format_atom_line`], and [`WriteGro::boxvecs`] need to be
/// specified. The default write implementations are based on these building blocks.
pub trait WriteGro<'a, A: 'a> {
    // Some generic getters that help build default implementations for `WriteGro::write`.
    /// Return the title for a `gro` file of this value.
    #[must_use]
    fn title(&self) -> String;

    /// Return the number of atoms in this value.
    ///
    /// # Invariant
    ///
    /// The returned number of atoms must be the same as the number of items yielded by
    /// [`WriteGro::atoms`]
    #[must_use]
    fn natoms(&self) -> usize;

    /// Return an iterator over `&A`.
    ///
    /// This `&A` item represents an atom to be formatted with [`WriteGro::format_atom_line`].
    ///
    /// # Invariant
    ///
    /// The returned iterator _must_ yield exactly the number of atoms specified by
    /// [`WriteGro::natoms`].
    #[must_use]
    fn atoms(&'a self) -> impl Iterator<Item = &'a A>;

    /// Format this value's box vectors for `gro`.
    ///
    /// # Note
    ///
    /// A valid box vector is either three _or_ nine space-separated floats.
    #[must_use]
    fn boxvecs(&self) -> String;

    // Format lines.
    /// Format a `&A` as a `gro` atom line.
    ///
    /// The line always has a trailing newline.
    #[must_use]
    fn format_atom_line(atom: &A) -> String;

    /// An iterator over `String`s representing this value's atoms as the atoms section of a `gro`
    /// file.
    ///
    /// Each line in the iterator has a trailing newline.
    #[must_use]
    fn format_atom_lines_iter(&'a self) -> impl Iterator<Item = String> {
        self.atoms().map(Self::format_atom_line)
    }

    /// Format this value's atoms as the atoms section of a `gro` file.
    ///
    /// The last line always has a trailing newline.
    #[must_use]
    fn format_atom_lines(&'a self) -> String {
        self.format_atom_lines_iter().collect()
    }

    // The actual write implementations.
    /// Write this value as `gro`.
    ///
    /// It is prudent to use this function over a [buffered writer](std::io::BufWriter).
    ///
    /// For a variant that uses multi-threaded formatting, see [`WriteGroPar::write_par`].
    ///
    /// # Examples
    ///
    /// [`Structure`] implements [`WriteGro`].
    ///
    /// ```
    /// use eightyseven::structure::{Atom, BoxVecs, Structure};
    /// use eightyseven::writer::WriteGro;
    /// use glam::Vec3;
    ///
    /// let structure = Structure {
    ///     title: "a simple structure".to_string(),
    ///     atoms: vec![
    ///         Atom {
    ///             resnum: 1,
    ///             resname: "MET".into(),
    ///             atomname: "CA".into(),
    ///             atomnum: 1,
    ///             position: Vec3::ZERO,
    ///             velocity: Vec3::ZERO,
    ///         }
    ///     ],
    ///     boxvecs: BoxVecs::Short([10.0, 10.0, 10.0]),
    /// };
    ///
    /// let writer = std::fs::File::create("simple.gro").unwrap();
    /// let mut writer = std::io::BufWriter::new(writer);
    /// structure.write(&mut writer).unwrap();
    /// ```
    fn write(&'a self, writer: &mut impl io::Write) -> io::Result<()> {
        // Write the header.
        writeln!(writer, "{}", self.title())?;
        writeln!(writer, "{}", self.natoms())?;

        // Format and write all of the lines.
        let lines = self.format_atom_lines_iter();
        for line in lines {
            writer.write_all(line.as_bytes())?;
        }

        // Write the box vectors.
        writeln!(writer, "{}", self.boxvecs())
    }

    /// Write this value to a `gro` file at a path.
    ///
    /// For a variant that uses multi-threaded formatting, see [`WriteGroPar::save_gro_par`].
    ///
    /// # Examples
    ///
    /// [`Structure`] implements [`WriteGro`].
    ///
    /// ```
    /// use eightyseven::structure::{Atom, BoxVecs, Structure};
    /// use eightyseven::writer::WriteGro;
    /// use glam::Vec3;
    ///
    /// let structure = Structure {
    ///     title: "a simple structure".to_string(),
    ///     atoms: vec![
    ///         Atom {
    ///             resnum: 1,
    ///             resname: "MET".into(),
    ///             atomname: "CA".into(),
    ///             atomnum: 1,
    ///             position: Vec3::ZERO,
    ///             velocity: Vec3::ZERO,
    ///         }
    ///     ],
    ///     boxvecs: BoxVecs::Short([10.0, 10.0, 10.0]),
    /// };
    ///
    /// let mut writer = std::fs::File::create("simple.gro").unwrap();
    /// structure.write(&mut writer).unwrap();
    /// ```
    fn save_gro<P: AsRef<Path>>(&'a self, path: P) -> io::Result<()> {
        let file = std::fs::File::create(path)?;
        let mut writer = std::io::BufWriter::new(file);
        self.write(&mut writer)
    }
}

/// A trait for writing out a value as `gro` using parallel formatting.
///
/// For single-threaded formatting, see [`WriteGro`].
///
/// `WriteGroPar` implements robust default implementations for writing to generic writers
/// ([`WriteGroPar::write_par`]) and to files ([`WriteGroPar::save_gro_par`]). To implement
/// this trait for a type, only [`WriteGroPar::atoms_par`] needs to be specified.
pub trait WriteGroPar<'a, A>: WriteGro<'a, A>
where
    A: 'a,
    &'a A: Send,
{
    // TODO: Can we really not just par_bridge WriteGro::atoms? The send bound is a requiredement
    // _here_, so if something is WriteGro<A> and we want to impl WriteGroPar for it, then that A
    // should be Send either way. ??
    #[must_use]
    fn atoms_par(&'a self) -> impl ParallelIterator<Item = &'a A>;

    /// Format this value's atoms as the atoms section of a `gro` file using parallel formatting.
    ///
    /// The last line always has a trailing newline.
    #[must_use]
    fn format_atom_lines_par(&'a self) -> String {
        self.atoms_par()
            .map(|atom| Self::format_atom_line(atom))
            .collect()
    }

    // The actual write implementations.
    /// Write this value as `gro` using parallel formatting.
    ///
    /// For single-threaded formatting, see [`WriteGro::write`].
    ///
    /// # Examples
    ///
    /// [`Structure`] implements [`WriteGroPar`].
    ///
    /// ```
    /// use eightyseven::structure::{Atom, BoxVecs, Structure};
    /// use eightyseven::writer::WriteGroPar;
    /// use glam::Vec3;
    ///
    /// let structure = Structure {
    ///     title: "a simple structure".to_string(),
    ///     atoms: vec![
    ///         Atom {
    ///             resnum: 1,
    ///             resname: "MET".into(),
    ///             atomname: "CA".into(),
    ///             atomnum: 1,
    ///             position: Vec3::ZERO,
    ///             velocity: Vec3::ZERO,
    ///         }
    ///     ],
    ///     boxvecs: BoxVecs::Short([10.0, 10.0, 10.0]),
    /// };
    ///
    /// let mut writer = std::fs::File::create("simple.gro").unwrap();
    /// structure.write_par(&mut writer).unwrap();
    /// ```
    fn write_par(&'a self, writer: &mut impl io::Write) -> io::Result<()> {
        // Write the header.
        writeln!(writer, "{}", self.title())?;
        writeln!(writer, "{}", self.natoms())?;

        // Format all of the lines.
        let lines = self.format_atom_lines_par();

        // Write the lines.
        write!(writer, "{lines}")?;

        // Write the box vectors.
        writeln!(writer, "{}", self.boxvecs())
    }

    /// Write this value to a `gro` file using parallel formatting.
    ///
    /// For single-threaded formatting, see [`WriteGro::save_gro`].
    ///
    /// # Examples
    ///
    /// [`Structure`] implements [`WriteGroPar`].
    ///
    /// ```
    /// use eightyseven::structure::{Atom, BoxVecs, Structure};
    /// use eightyseven::writer::WriteGroPar;
    /// use glam::Vec3;
    ///
    /// let structure = Structure {
    ///     title: "a simple structure".to_string(),
    ///     atoms: vec![
    ///         Atom {
    ///             resnum: 1,
    ///             resname: "MET".into(),
    ///             atomname: "CA".into(),
    ///             atomnum: 1,
    ///             position: Vec3::ZERO,
    ///             velocity: Vec3::ZERO,
    ///         }
    ///     ],
    ///     boxvecs: BoxVecs::Short([10.0, 10.0, 10.0]),
    /// };
    ///
    /// structure.save_gro_par("simple.gro").unwrap();
    /// ```
    fn save_gro_par<P: AsRef<Path>>(&'a self, path: P) -> io::Result<()> {
        let mut file = std::fs::File::create(path)?;
        self.write_par(&mut file)
    }
}

impl<'a> WriteGro<'a, Atom> for Structure {
    fn title(&self) -> String {
        self.title.clone()
    }

    fn natoms(&self) -> usize {
        self.natoms()
    }

    fn atoms(&'a self) -> impl Iterator<Item = &'a Atom> {
        self.atoms.iter()
    }

    fn boxvecs(&self) -> String {
        self.boxvecs.to_string()
    }

    fn format_atom_line(atom: &Atom) -> String {
        format_atom_line(
            atom.resnum,
            atom.resname,
            atom.atomname,
            atom.atomnum,
            atom.position.to_array(),
            match atom.velocity.to_array() {
                [0.0, 0.0, 0.0] => None,
                vel => Some(vel),
            },
        )
    }
}

impl<'a> WriteGroPar<'a, Atom> for Structure {
    fn atoms_par(&'a self) -> impl ParallelIterator<Item = &'a Atom> {
        self.atoms.par_iter()
    }
}

/// Format a single [`Atom`] as a `gro` atom line.
///
/// The returned line always has a trailing newline.
///
/// Note that `resname` and `atomname` are assumed have a `len` of at most 5 characters. These
/// fields in the `gro` line are both exactly 5 columns wide, hence the requirement. This condition
/// is asserted.
///
/// # Examples
///
/// ```
/// use glam::Vec3;
/// use eightyseven::{structure::Atom, writer::format_atom_line};
///
/// let resnum = 1;
/// let resname = "MET";
/// let atomname = "CA";
/// let atomnum = 1;
/// let position = Vec3::X;
/// let velocity = Vec3::Y;
///
/// let line = "    1MET     CA    1   1.000   0.000   0.000  0.0000  1.0000  0.0000\n";
/// let formatted = format_atom_line(
///     resnum,
///     resname,
///     atomname,
///     atomnum,
///     position.to_array(),
///     Some(velocity.to_array())
/// );
/// assert_eq!(formatted, line);
/// ```
///
/// # Panics
///
/// This function will panic if `resname` or `atomname` has a `len` greater than 5.
#[must_use]
pub fn format_atom_line(
    resnum: u32,
    // TODO: resname and atomname can only be <= 5 chars long. Make sure that this is understood on the type level. Or even that it is ascii, and that will invite some more cool optimizations?
    resname: impl ToString,
    atomname: impl ToString,
    atomnum: u32,
    position: [f32; 3],
    velocity: Option<[f32; 3]>,
) -> String {
    // Wrap atomnum and resnum, since they should not be more than 5 characters wide.
    let atomnum = atomnum % GRO_INTEGER_LIMIT;
    let resnum = resnum % GRO_INTEGER_LIMIT;
    // The `to_string` is necessary, here, since the `:>5` alignment is otherwise ignored.
    let resname = resname.to_string();
    let atomname = atomname.to_string();

    assert!(resname.len() <= 5);
    assert!(atomname.len() <= 5);

    let [x, y, z] = position;
    if let Some([vx, vy, vz]) = velocity {
        format!("{resnum:>5}{resname:<5}{atomname:>5}{atomnum:>5}{x:8.3}{y:8.3}{z:8.3}{vx:8.4}{vy:8.4}{vz:8.4}\n")
    } else {
        format!("{resnum:>5}{resname:<5}{atomname:>5}{atomnum:>5}{x:8.3}{y:8.3}{z:8.3}\n")
    }
}

#[cfg(test)]
mod tests {
    use std::io::{self, Read};

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

    #[test]
    fn in_out() -> io::Result<()> {
        // Read in the data.
        let mut gro = Vec::new();
        std::fs::File::open(crate::tests::PATH)?.read_to_end(&mut gro)?;
        let structure = Structure::read(gro.as_slice())?;

        // Write it out.
        let mut out = Vec::new();
        structure.write(&mut out)?;

        // Read in that output for the comparison.
        let processed = Structure::read(out.as_slice())?;

        // Check that they are identical.
        assert_eq!(structure.title, processed.title);
        assert_eq!(structure.natoms(), processed.natoms());
        assert_eq!(structure.boxvecs, processed.boxvecs);
        for idx in 0..structure.natoms() {
            let original = structure.atoms[idx];
            let processed = processed.atoms[idx];
            assert_eq!(original, processed, "atom {idx} does not match");
        }

        Ok(())
    }

    #[test]
    fn in_out_par() -> io::Result<()> {
        // Read in the data.
        let structure = Structure::open_gro(crate::tests::PATH)?;

        // Write it out.
        let mut out_par = Vec::new();
        structure.write_par(&mut out_par)?;

        // Read in that output for the comparison.
        let processed_par = Structure::read(out_par.as_slice())?;

        // Check that they are identical.
        assert_eq!(structure.title, processed_par.title);
        assert_eq!(structure.natoms(), processed_par.natoms());
        assert_eq!(structure.boxvecs, processed_par.boxvecs);
        for idx in 0..structure.natoms() {
            let original = structure.atoms[idx];
            let processed_par = processed_par.atoms[idx];
            assert_eq!(original, processed_par, "atom {idx} does not match");
        }

        Ok(())
    }

    #[test]
    fn seq_and_par() -> io::Result<()> {
        // Read in the data.
        let structure = Structure::open_gro(crate::tests::PATH)?;

        // Write it out.
        let mut out = Vec::new();
        structure.write(&mut out)?;
        let mut out_par = Vec::new();
        structure.write_par(&mut out_par)?;

        // Read in that output for the comparison.
        let processed = Structure::read(out.as_slice())?;
        let processed_par = Structure::read(out_par.as_slice())?;

        // Check that they are identical.
        assert_eq!(processed.title, processed_par.title);
        assert_eq!(processed.natoms(), processed_par.natoms());
        assert_eq!(processed.boxvecs, processed_par.boxvecs);
        for idx in 0..structure.natoms() {
            let processed = processed.atoms[idx];
            let processed_par = processed_par.atoms[idx];
            assert_eq!(processed, processed_par, "atom {idx} does not match");
        }

        Ok(())
    }
}