eightyseven 0.1.5

use std::io::BufReader;
use std::path::Path;
use std::{io, str::FromStr};

use glam::Vec3;

use crate::structure::{Atom, AtomName, AtomNum, BoxVecs, ResName, ResNum, Structure};

/// The error type for reading `gro` files.
///
/// A distinction is made between _I/O_, _unexpected end-of_, and _bad-value_ errors.
///
/// A [`ParseGroError`] can always be cast to an [`io::Error`]. If the error is of the
/// [`ParseGroError::IOError`] variant, the inner `io::Error` will be used directly. Otherwise, the
/// error is converted to an instance of [`io::ErrorKind::Other`].
///
/// Similarly, any `io::Error` can be cast to a `ParseGroError`. This makes conversion between
/// error types convenient and allows for using the try-operator (`?`) in both
/// [`reader::Result`][`Result`] and [`io::Result`]-returning function.
#[non_exhaustive]
#[derive(Debug)]
pub enum ParseGroError {
    IOError(io::Error),

    UnexpectedEOF(ExpectedItem),
    UnexpectedEOL(ExpectedItem),

    BadNAtoms(std::num::ParseIntError),
    InsufficientAtoms { exp: usize, enc: usize },
    BadResnum(std::num::ParseIntError),
    BadAtomnum(std::num::ParseIntError),
    BadPositionValue(std::num::ParseFloatError),
    BadVelocityValue(std::num::ParseFloatError),
    BadBoxVecsLength(usize),
    BadBoxVecsValue(std::num::ParseFloatError),
}

#[non_exhaustive]
#[derive(Debug, Clone, Copy)]
/// The expected values for both the [`ParseGroError::UnexpectedEOF`] and
/// [`ParseGroError::UnexpectedEOL`] variants.
pub enum ExpectedItem {
    // UnexpectedEOF.
    Title,   // "title"
    NAtoms,  // "number of atoms"
    BoxVecs, // "box vectors"

    // UnexpectedEOL.
    Resnum,   // "resnum"
    Resname,  // "resname"
    Atomname, // "atomname"
    Atomnum,  // "atomnum"
    Position, // "position"
    Velocity, // "velocity"
}

impl std::fmt::Display for ParseGroError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::IOError(err) => write!(f, "I/O error: {err:?}"),

            Self::UnexpectedEOF(expected) => {
                write!(f, "unexpected end of file: expected {expected}")
            }
            Self::UnexpectedEOL(expected) => {
                write!(f, "unexpected end of line: while reading {expected}")
            }

            Self::BadNAtoms(err) => write!(f, "could not read the number of atoms: {err}"),
            Self::InsufficientAtoms { exp, enc } => write!(
                f,
                "could not read the specified number of atoms: \
                    expected {exp} but could only read {enc}"
            ),
            Self::BadResnum(err) => write!(f, "could not read the residue number: {err}"),
            Self::BadAtomnum(err) => write!(f, "could not read the atom number: {err}"),
            Self::BadPositionValue(err) => write!(f, "could not read position value: {err}"),
            Self::BadVelocityValue(err) => write!(f, "could not read velocity value: {err}"),
            Self::BadBoxVecsLength(n) => {
                write!(
                    f,
                    "could not read box vectors: expected 3 or 9 values, found {n}"
                )
            }
            Self::BadBoxVecsValue(err) => write!(f, "could not read box vectors value: {err}"),
        }
    }
}

impl std::error::Error for ParseGroError {}

impl std::fmt::Display for ExpectedItem {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Self::Title => "title",
            Self::NAtoms => "number of atoms",
            Self::BoxVecs => "box vectors",
            Self::Resnum => "resnum",
            Self::Resname => "resname",
            Self::Atomname => "atomname",
            Self::Atomnum => "atomnum",
            Self::Position => "position",
            Self::Velocity => "velocity",
        })
    }
}

impl From<ParseGroError> for io::Error {
    fn from(err: ParseGroError) -> Self {
        match err {
            // Just give the inner io::Error.
            ParseGroError::IOError(err) => err,
            // Wrap them as an io::ErrorKind::Other.
            other => io::Error::other(other),
        }
    }
}

impl From<io::Error> for ParseGroError {
    fn from(err: io::Error) -> Self {
        ParseGroError::IOError(err)
    }
}

pub type Result<T> = std::result::Result<T, ParseGroError>;

/// The ranges within a single atom line of a `gro` file.
pub mod ranges {
    use std::ops::Range;

    pub const RESNUM: Range<usize> = 0..5;
    pub const RESNAME: Range<usize> = 5..10;
    pub const ATOMNAME: Range<usize> = 10..15;
    pub const ATOMNUM: Range<usize> = 15..20;

    // NOTE: These values only hold when the default `f8.3` and `f8.4` precisions are used. This
    // may change in the future when I implement non-standard precisions.
    pub const POSITION: Range<usize> = 20..44;
    pub const VELOCITY: Range<usize> = 44..68;

    /// Fields at least up to and including `position` should be there for a valid `gro` atom line.
    pub const MIN_LINE_LEN: usize = POSITION.end;
}

/// Specify which fields must be read and which must be skipped while reading a `gro` atom line.
///
/// [`ParseList::ALL`] is considered the default configuration.
///
/// If a value is set to `true` (as set by [`ParseList::ALL`]), the `parse_*` for that value will
/// be called and the result will be presented as `Some(value)` to [`ReadGro<A>::build_atom`].
pub struct ParseList {
    pub resnum: bool,
    pub resname: bool,
    pub atomname: bool,
    pub atomnum: bool,
    pub position: bool,
    pub velocity: bool,
}

impl ParseList {
    pub const ALL: Self = Self {
        resnum: true,
        resname: true,
        atomname: true,
        atomnum: true,
        position: true,
        velocity: true,
    };
}

/// A trait for reading `gro` files int.
///
/// `ReadGro` implements robust default implementations for reading from generic readers
/// ([`ReadGro::read`]), and to files ([`ReadGro::read_from_file`]) and
/// ([`ReadGro::open_gro`]).
///
/// To implement this trait for a type, only [`ReadGro::build_atom`] and
/// [`ReadGro::build_structure`] need to be specified. The default read implementations are based
/// on these building blocks.
///
/// It is possible to read only the parts from the atom lines that are necessary for the atom (`A`)
/// type with the [`ReadGro::PARSE_LIST`] const struct. For example, if you are only interested in
/// the positions, the preceding fields can be skipped and only the positions field will be read.
/// Care must be taken to set up a good match between the implementation of `ReadGro::build_atom`
/// and the `ReadGro::PARSE_LIST`. The implementation of `ReadGro` for [`Structure`] is a good
/// example of this.
pub trait ReadGro<A>: Sized {
    const PARSE_LIST: ParseList = ParseList::ALL;

    /// Build an atom (`A`) from options of components.
    ///
    /// Together with setting the [`ReadGro<A>::PARSE_LIST`], it is possible to read only the
    /// values that are necessary to construct `A`.
    ///
    /// # Note
    ///
    /// It is very advisable to mark the implementation for this function as `#[inline]`. This
    /// function will be called in a rather hot loop, internally.
    fn build_atom(
        resnum: Option<ResNum>,
        resname: Option<ResName>,
        atomname: Option<AtomName>,
        atomnum: Option<AtomNum>,
        position: Option<[f32; 3]>,
        velocity: Option<[f32; 3]>,
    ) -> A;

    /// Build the value from a title, atoms, and box vectors.
    fn build_structure(title: String, atoms: Vec<A>, boxvecs: BoxVecs) -> Self;

    #[inline]
    fn parse_resnum(line: &str) -> Result<ResNum> {
        line.get(ranges::RESNUM)
            .ok_or(ParseGroError::UnexpectedEOL(ExpectedItem::Resnum))?
            .trim()
            .parse()
            .map_err(ParseGroError::BadResnum)
    }

    #[inline]
    fn parse_resname(line: &str) -> Result<ResName> {
        line.get(ranges::RESNAME)
            .ok_or(ParseGroError::UnexpectedEOL(ExpectedItem::Resname))
            .map(|s| s.trim().into())
    }

    #[inline]
    fn parse_atomname(line: &str) -> Result<AtomName> {
        line.get(ranges::ATOMNAME)
            .ok_or(ParseGroError::UnexpectedEOL(ExpectedItem::Atomname))
            .map(|s| s.trim().into())
    }

    #[inline]
    fn parse_atomnum(line: &str) -> Result<AtomNum> {
        line.get(ranges::ATOMNUM)
            .ok_or(ParseGroError::UnexpectedEOL(ExpectedItem::Atomnum))?
            .trim()
            .parse()
            .map_err(ParseGroError::BadAtomnum)
    }

    #[inline]
    fn parse_position(line: &str) -> Result<[f32; 3]> {
        parse_floats(
            transpose_option_array([20..28, 28..36, 36..44].map(|r| line.get(r)))
                .ok_or(ParseGroError::UnexpectedEOL(ExpectedItem::Position))?
                .map(str::trim),
        )
        .map_err(ParseGroError::BadPositionValue)
    }

    #[inline]
    fn parse_velocity(line: &str) -> Result<[f32; 3]> {
        parse_floats(
            transpose_option_array([44..52, 52..60, 60..68].map(|r| line.get(r)))
                .ok_or(ParseGroError::UnexpectedEOL(ExpectedItem::Velocity))?
                .map(str::trim),
        )
        .map_err(ParseGroError::BadVelocityValue)
    }

    /// Parse an atom line.
    ///
    /// # Examples
    ///
    /// [`ReadGro`] is implement for [`Structure`].
    ///
    /// ```
    /// use eightyseven::reader::ReadGro;
    /// use eightyseven::structure::Structure;
    ///
    /// let line = "    1MET     BB    1   7.508   4.691   2.177  0.0306 -0.1635  0.1420";
    /// let atom = Structure::parse_atom_line(line).unwrap();
    ///
    /// assert_eq!(atom.resnum, 1);
    /// assert_eq!(atom.resname.as_str(), "MET");
    /// assert_eq!(atom.atomname.as_str(), "BB");
    /// assert_eq!(atom.atomnum, 1);
    /// ```
    #[rustfmt::skip]
    #[inline]
    fn parse_atom_line(line: &str) -> Result<A> {
        assert!(line.len() >= ranges::MIN_LINE_LEN);
        let resnum =   if Self::PARSE_LIST.resnum   { Some(Self::parse_resnum(line)?)   } else { None };
        let resname =  if Self::PARSE_LIST.resname  { Some(Self::parse_resname(line)?)  } else { None };
        let atomname = if Self::PARSE_LIST.atomname { Some(Self::parse_atomname(line)?) } else { None };
        let atomnum =  if Self::PARSE_LIST.atomnum  { Some(Self::parse_atomnum(line)?)  } else { None };
        let position = if Self::PARSE_LIST.position { Some(Self::parse_position(line)?) } else { None };
        let velocity = if Self::PARSE_LIST.velocity {
            if line.len() > 44 {
                Some(Self::parse_velocity(line)?)
            } else {
                None
            }
        } else {
            None
        };

        Ok(Self::build_atom(
            resnum, resname, atomname, atomnum, position, velocity,
        ))
    }

    /// Parse a box vectors line.
    ///
    /// # Examples
    ///
    /// [`ReadGro`] is implement for [`Structure`].
    ///
    /// ```
    /// use eightyseven::reader::ReadGro;
    /// use eightyseven::structure::{BoxVecs, Structure};
    ///
    /// let short = "123.0 456.0 789.0";
    /// let full = "0.0 1.0 2.0 3.0 4.0 5.0 6.0 8.0 9.0";
    ///
    /// assert_eq!(Structure::parse_boxvecs(short).unwrap(), BoxVecs::Short([123.0, 456.0, 789.0]));
    /// assert_eq!(Structure::parse_boxvecs(full).unwrap(), BoxVecs::Full([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 8.0, 9.0]));
    ///
    /// let broken = "51.72 51.12";
    ///
    /// assert!(Structure::parse_boxvecs(broken).is_err());
    /// ```
    fn parse_boxvecs(line: &str) -> Result<BoxVecs> {
        let vs: Vec<&str> = line.split_ascii_whitespace().collect();
        match vs.len() {
            // The try_into().unwrap()s here are safe, since we just checked the len.
            3 => Ok(BoxVecs::Short(
                parse_floats(vs.try_into().unwrap()).map_err(ParseGroError::BadBoxVecsValue)?,
            )),
            9 => Ok(BoxVecs::Full(
                parse_floats(vs.try_into().unwrap()).map_err(ParseGroError::BadBoxVecsValue)?,
            )),
            n => Err(ParseGroError::BadBoxVecsLength(n)),
        }
    }

    /// Read `gro` data.
    ///
    /// # Examples
    ///
    /// [`ReadGro`] is implement for [`Structure`].
    ///
    /// ```
    /// use std::io::{BufReader, Read};
    ///
    /// use eightyseven::reader::ReadGro;
    /// use eightyseven::structure::Structure;
    ///
    /// let mut buf = Vec::new();
    /// std::fs::File::open("tests/eq.gro").unwrap().read_to_end(&mut buf).unwrap();
    /// let reader = BufReader::new(buf.as_slice());
    /// let structure = Structure::read(reader).expect("should be a valid gro file");
    /// println!("{}", structure.title);
    /// ```
    fn read(reader: impl io::BufRead) -> Result<Self> {
        // We will iterate over lines.
        let mut lines = reader.lines();

        // Read the header.
        let title = String::from(
            lines
                .next()
                .ok_or(ParseGroError::UnexpectedEOF(ExpectedItem::Title))??
                .trim(),
        );
        let natoms = lines
            .next()
            .ok_or(ParseGroError::UnexpectedEOF(ExpectedItem::NAtoms))??
            .trim()
            .parse()
            .map_err(ParseGroError::BadNAtoms)?;

        // Read all of the atoms.
        let mut atoms = Vec::with_capacity(natoms);
        for idx in 0..natoms {
            let atom = Self::parse_atom_line(&lines.next().ok_or(
                ParseGroError::InsufficientAtoms {
                    exp: natoms,
                    enc: idx,
                },
            )??)?;
            atoms.push(atom);
        }

        // Finally, the box vectors.
        let boxvecs = Self::parse_boxvecs(
            &lines
                .next()
                .ok_or(ParseGroError::UnexpectedEOF(ExpectedItem::BoxVecs))??,
        )?;

        Ok(Self::build_structure(title, atoms, boxvecs))
    }

    /// Read a `gro` file from a [`std::fs::File`].
    ///
    /// # Examples
    ///
    /// [`ReadGro`] is implement for [`Structure`].
    ///
    /// ```
    /// use eightyseven::reader::ReadGro;
    /// use eightyseven::structure::Structure;
    ///
    /// let file = std::fs::File::open("tests/eq.gro").unwrap();
    /// let structure = Structure::read_from_file(file).expect("should be a valid gro file");
    /// println!("{}", structure.title);
    /// ```
    fn read_from_file(file: std::fs::File) -> Result<Self> {
        let reader = BufReader::new(file);
        Self::read(reader)
    }

    /// Read a `gro` file from a file at a path.
    ///
    /// # Examples
    ///
    /// [`ReadGro`] is implement for [`Structure`].
    ///
    /// ```
    /// use eightyseven::reader::ReadGro;
    /// use eightyseven::structure::Structure;
    ///
    /// let structure = Structure::open_gro("tests/eq.gro").expect("should be a valid gro file");
    /// println!("{}", structure.title);
    /// ```
    fn open_gro<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file = std::fs::File::open(path)?;
        Self::read_from_file(file)
    }
}

impl ReadGro<Atom> for Structure {
    #[inline]
    fn build_atom(
        resnum: Option<ResNum>,
        resname: Option<ResName>,
        atomname: Option<AtomName>,
        atomnum: Option<AtomNum>,
        position: Option<[f32; 3]>,
        velocity: Option<[f32; 3]>,
    ) -> Atom {
        // We know that the values we unwrap here should be `Some`. Any errors were handled
        // upstream and we use the default `ParseList::All` configuration.
        Atom {
            resnum: resnum.unwrap(),
            resname: resname.unwrap(),
            atomname: atomname.unwrap(),
            atomnum: atomnum.unwrap(),
            position: Vec3::from_array(position.unwrap()),
            velocity: Vec3::from_array(velocity.unwrap_or_default()),
        }
    }

    fn build_structure(title: String, atoms: Vec<Atom>, boxvecs: BoxVecs) -> Self {
        Self {
            title,
            atoms,
            boxvecs,
        }
    }
}

impl FromStr for Structure {
    type Err = ParseGroError;

    fn from_str(gro: &str) -> Result<Self> {
        Self::read(gro.as_bytes())
    }
}

/// Transpose an array of options into an option of an array.
///
/// If any of the values is [`None`], this function will return [`None`].
/// If all values are [`Some`], an array of the unwrapped values is returned.
fn transpose_option_array<const N: usize, T>(vs: [Option<T>; N]) -> Option<[T; N]> {
    if vs.iter().any(Option::is_none) {
        None
    } else {
        Some(vs.map(Option::unwrap))
    }
}

/// Parse an array of `&str`s into `f32`s.
#[inline]
fn parse_floats<const N: usize>(
    vs: [&str; N],
) -> std::result::Result<[f32; N], std::num::ParseFloatError> {
    vs.iter()
        .map(|v| v.parse())
        .collect::<std::result::Result<Vec<_>, _>>()
        .map(|vs| vs.try_into().unwrap()) // Safe, since we know N_in == N_out.
}

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

    use super::*;

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

    #[test]
    fn open_gro() -> io::Result<()> {
        let structure = Structure::open_gro(crate::tests::PATH)?;
        assert_eq!(structure.title, "cg protein in water");
        assert_eq!(structure.natoms(), 2869);

        let center = Vec3::new(3.9875, 3.9760, 2.7035);
        assert!(structure.center().abs_diff_eq(center, EPS));
        assert_eq!(
            structure.atoms[123].position,
            Vec3::new(5.589, 5.256, 3.264)
        );
        assert_eq!(
            structure.atoms[456].velocity,
            Vec3::new(-0.144, -0.2159, -0.1276)
        );

        assert_eq!(
            structure.boxvecs,
            BoxVecs::Full([7.59661, 7.59661, 5.37162, 0.0, 0.0, 0.0, 0.0, 3.79831, 3.79831])
        );

        Ok(())
    }
}