molcrafts-molrs 0.7.0

//! Tripos MOL2 reader and writer.
//!
//! MOL2 is a section-delimited text format used widely in cheminformatics.
//! Sections are introduced by lines starting with `@<TRIPOS>`. The four
//! sections this reader honours:
//!
//! - `MOLECULE` — name + counts line (`n_atoms n_bonds n_subst n_feat n_sets`)
//! - `ATOM` — `id name x y z atom_type [subst_id subst_name [charge]]`
//! - `BOND` — `id atom_i atom_j type` (type ∈ `1 2 3 am ar du un nc`)
//! - All other sections are tolerantly skipped.
//!
//! Multi-molecule files (repeated `@<TRIPOS>MOLECULE` blocks) are supported by
//! repeated [`Mol2Reader::read_frame`] calls.
//!
//! ## Output Frame
//!
//! - `"atoms"` block: `id` (i32, 1-based), `name` (str), `x`/`y`/`z` (F),
//!   `atom_type` (str), `subst_id` (i32) and `subst_name` (str) when present,
//!   `charge` (F) when present.
//! - `"bonds"` block (when present): `atomi`/`atomj` (u32, 0-based indices),
//!   `bond_type` (str).
//! - `frame.meta["title"]` = molecule name.

use std::io::{BufRead, BufWriter, Error, ErrorKind, Result, Write};
use std::path::Path;

use ndarray::{Array1, IxDyn};

use molrs::store::block::Block;
use molrs::store::frame::Frame;
use molrs::types::{F, I, U};

use crate::io::reader::{FrameReader, Reader};
use crate::io::writer::{FrameWriter, Writer};

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn invalid_data<E: std::fmt::Display>(e: E) -> Error {
    Error::new(ErrorKind::InvalidData, e.to_string())
}

fn insert_float_col(block: &mut Block, key: &str, vals: Vec<F>) -> Result<()> {
    let n = vals.len();
    let arr = Array1::from_vec(vals)
        .into_shape_with_order(IxDyn(&[n]))
        .map_err(invalid_data)?
        .into_dyn();
    block.insert(key, arr).map_err(invalid_data)
}

fn insert_int_col(block: &mut Block, key: &str, vals: Vec<I>) -> Result<()> {
    let n = vals.len();
    let arr = Array1::from_vec(vals)
        .into_shape_with_order(IxDyn(&[n]))
        .map_err(invalid_data)?
        .into_dyn();
    block.insert(key, arr).map_err(invalid_data)
}

fn insert_uint_col(block: &mut Block, key: &str, vals: Vec<U>) -> Result<()> {
    let n = vals.len();
    let arr = Array1::from_vec(vals)
        .into_shape_with_order(IxDyn(&[n]))
        .map_err(invalid_data)?
        .into_dyn();
    block.insert(key, arr).map_err(invalid_data)
}

fn insert_str_col(block: &mut Block, key: &str, vals: Vec<String>) -> Result<()> {
    let n = vals.len();
    let arr = Array1::from_vec(vals)
        .into_shape_with_order(IxDyn(&[n]))
        .map_err(invalid_data)?
        .into_dyn();
    block.insert(key, arr).map_err(invalid_data)
}

// ---------------------------------------------------------------------------
// Parsed records
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
struct Mol2Atom {
    id: I,
    name: String,
    x: F,
    y: F,
    z: F,
    atom_type: String,
    subst_id: Option<I>,
    subst_name: Option<String>,
    charge: Option<F>,
}

fn parse_atom(line: &str, line_no: usize) -> Result<Mol2Atom> {
    let toks: Vec<&str> = line.split_whitespace().collect();
    if toks.len() < 6 {
        return Err(invalid_data(format!(
            "line {}: ATOM row needs ≥6 fields, got {}",
            line_no,
            toks.len()
        )));
    }
    let id: I = toks[0]
        .parse()
        .map_err(|_| invalid_data(format!("line {}: bad atom id '{}'", line_no, toks[0])))?;
    let name = toks[1].to_string();
    let x: F = toks[2]
        .parse()
        .map_err(|_| invalid_data(format!("line {}: bad x '{}'", line_no, toks[2])))?;
    let y: F = toks[3]
        .parse()
        .map_err(|_| invalid_data(format!("line {}: bad y '{}'", line_no, toks[3])))?;
    let z: F = toks[4]
        .parse()
        .map_err(|_| invalid_data(format!("line {}: bad z '{}'", line_no, toks[4])))?;
    let atom_type = toks[5].to_string();
    let subst_id = toks.get(6).and_then(|t| t.parse::<I>().ok());
    let subst_name = toks.get(7).map(|t| t.to_string());
    let charge = toks.get(8).and_then(|t| t.parse::<F>().ok());
    Ok(Mol2Atom {
        id,
        name,
        x,
        y,
        z,
        atom_type,
        subst_id,
        subst_name,
        charge,
    })
}

#[derive(Debug, Clone)]
struct Mol2Bond {
    _id: I,
    atom_i: I,
    atom_j: I,
    bond_type: String,
}

fn parse_bond(line: &str, line_no: usize) -> Result<Mol2Bond> {
    let toks: Vec<&str> = line.split_whitespace().collect();
    if toks.len() < 4 {
        return Err(invalid_data(format!(
            "line {}: BOND row needs ≥4 fields, got {}",
            line_no,
            toks.len()
        )));
    }
    Ok(Mol2Bond {
        _id: toks[0].parse().unwrap_or(0),
        atom_i: toks[1]
            .parse()
            .map_err(|_| invalid_data(format!("line {}: bad atom_i '{}'", line_no, toks[1])))?,
        atom_j: toks[2]
            .parse()
            .map_err(|_| invalid_data(format!("line {}: bad atom_j '{}'", line_no, toks[2])))?,
        bond_type: toks[3].to_string(),
    })
}

// ---------------------------------------------------------------------------
// Reader
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq)]
enum Section {
    None,
    Molecule,
    Atom,
    Bond,
    Other,
}

fn section_for(line: &str) -> Section {
    let trimmed = line.trim_end();
    let upper = trimmed.trim_start_matches("@<TRIPOS>");
    if !trimmed.starts_with("@<TRIPOS>") {
        return Section::None;
    }
    match upper {
        "MOLECULE" => Section::Molecule,
        "ATOM" => Section::Atom,
        "BOND" => Section::Bond,
        _ => Section::Other,
    }
}

/// State machine for reading a single MOL2 record.
fn read_one_record<R: BufRead>(
    reader: &mut R,
    pending_first_line: &mut Option<String>,
) -> Result<Option<Frame>> {
    // Find the start of a MOLECULE block.
    let mut buf = String::new();

    loop {
        if let Some(line) = pending_first_line.take() {
            if section_for(&line) == Section::Molecule {
                break;
            }
            // Otherwise discard and keep scanning.
            continue;
        }
        buf.clear();
        let bytes = reader.read_line(&mut buf)?;
        if bytes == 0 {
            return Ok(None);
        }
        if section_for(&buf) == Section::Molecule {
            break;
        }
    }

    // ---- MOLECULE block ----
    // Line 1: name
    buf.clear();
    if reader.read_line(&mut buf)? == 0 {
        return Err(invalid_data("EOF after @<TRIPOS>MOLECULE"));
    }
    let name = buf.trim().to_string();

    // Line 2: counts
    buf.clear();
    if reader.read_line(&mut buf)? == 0 {
        return Err(invalid_data("EOF before MOLECULE counts"));
    }
    let counts: Vec<&str> = buf.split_whitespace().collect();
    if counts.is_empty() {
        return Err(invalid_data("empty MOLECULE counts line"));
    }
    let n_atoms: usize = counts[0]
        .parse()
        .map_err(|_| invalid_data(format!("bad MOLECULE atom count '{}'", counts[0])))?;
    let n_bonds: usize = counts.get(1).and_then(|s| s.parse().ok()).unwrap_or(0);

    // The next several lines (mol_type, charge_type, [status_bits], [comment])
    // are optional and we don't need them. Skip until we hit the next section
    // marker or a blank line followed by a section marker.
    let mut atoms: Vec<Mol2Atom> = Vec::with_capacity(n_atoms);
    let mut bonds: Vec<Mol2Bond> = Vec::with_capacity(n_bonds);
    let mut current = Section::Molecule;
    let mut atoms_remaining = n_atoms;
    let mut bonds_remaining = n_bonds;

    let mut line_no = 2usize;
    loop {
        buf.clear();
        let bytes = reader.read_line(&mut buf)?;
        if bytes == 0 {
            // EOF — only valid if we collected enough atoms/bonds
            break;
        }
        line_no += 1;

        // Section change?
        let sec = section_for(&buf);
        match sec {
            Section::None => {}
            Section::Molecule => {
                // Next molecule begins. Push back this line so the next
                // read_one_record sees it.
                *pending_first_line = Some(buf.clone());
                break;
            }
            other => {
                current = other;
                continue;
            }
        }

        let trimmed = buf.trim();
        if trimmed.is_empty() || trimmed.starts_with('#') {
            continue;
        }

        match current {
            Section::Atom if atoms_remaining > 0 => {
                atoms.push(parse_atom(&buf, line_no)?);
                atoms_remaining -= 1;
            }
            Section::Bond if bonds_remaining > 0 => {
                bonds.push(parse_bond(&buf, line_no)?);
                bonds_remaining -= 1;
            }
            _ => {} // tolerate / skip
        }
    }

    if atoms.is_empty() {
        return Err(invalid_data(format!(
            "MOLECULE '{}' yielded zero atoms (declared {})",
            name, n_atoms
        )));
    }

    Ok(Some(build_frame(name, atoms, bonds)?))
}

fn build_frame(name: String, atoms: Vec<Mol2Atom>, bonds: Vec<Mol2Bond>) -> Result<Frame> {
    let n = atoms.len();
    let mut block = Block::new();
    let mut id = Vec::with_capacity(n);
    let mut a_name = Vec::with_capacity(n);
    let mut x = Vec::with_capacity(n);
    let mut y = Vec::with_capacity(n);
    let mut z = Vec::with_capacity(n);
    let mut a_type = Vec::with_capacity(n);
    let mut subst_id = Vec::with_capacity(n);
    let mut subst_name = Vec::with_capacity(n);
    let mut charge = Vec::with_capacity(n);
    let mut have_subst = false;
    let mut have_charge = false;
    for a in &atoms {
        id.push(a.id);
        a_name.push(a.name.clone());
        x.push(a.x);
        y.push(a.y);
        z.push(a.z);
        a_type.push(a.atom_type.clone());
        subst_id.push(a.subst_id.unwrap_or(0));
        subst_name.push(a.subst_name.clone().unwrap_or_default());
        charge.push(a.charge.unwrap_or(0.0));
        if a.subst_id.is_some() {
            have_subst = true;
        }
        if a.charge.is_some() {
            have_charge = true;
        }
    }
    insert_int_col(&mut block, "id", id)?;
    insert_str_col(&mut block, "name", a_name)?;
    insert_float_col(&mut block, "x", x)?;
    insert_float_col(&mut block, "y", y)?;
    insert_float_col(&mut block, "z", z)?;
    insert_str_col(&mut block, "atom_type", a_type)?;
    if have_subst {
        insert_int_col(&mut block, "subst_id", subst_id)?;
        insert_str_col(&mut block, "subst_name", subst_name)?;
    }
    if have_charge {
        insert_float_col(&mut block, "charge", charge)?;
    }

    let mut frame = Frame::new();
    if !name.is_empty() {
        frame.meta.insert("title".into(), name);
    }
    frame.insert("atoms", block);

    if !bonds.is_empty() {
        let bn = bonds.len();
        let mut atomi = Vec::with_capacity(bn);
        let mut atomj = Vec::with_capacity(bn);
        let mut btype = Vec::with_capacity(bn);
        for b in &bonds {
            if b.atom_i <= 0 || b.atom_j <= 0 {
                return Err(invalid_data(format!(
                    "non-positive atom index in bond: {} - {}",
                    b.atom_i, b.atom_j
                )));
            }
            if (b.atom_i as usize) > n || (b.atom_j as usize) > n {
                return Err(invalid_data(format!(
                    "bond references out-of-range atom: {} - {} (have {} atoms)",
                    b.atom_i, b.atom_j, n
                )));
            }
            atomi.push((b.atom_i as U) - 1);
            atomj.push((b.atom_j as U) - 1);
            btype.push(b.bond_type.clone());
        }
        let mut bblock = Block::new();
        insert_uint_col(&mut bblock, "atomi", atomi)?;
        insert_uint_col(&mut bblock, "atomj", atomj)?;
        insert_str_col(&mut bblock, "bond_type", btype)?;
        frame.insert("bonds", bblock);
    }

    Ok(frame)
}

/// Public single-frame helper for callers that just want the first molecule.
pub fn read_mol2<P: AsRef<Path>>(path: P) -> Result<Frame> {
    let file = std::fs::File::open(path.as_ref())?;
    let mut reader = std::io::BufReader::new(file);
    let mut pending = None;
    read_one_record(&mut reader, &mut pending)?
        .ok_or_else(|| invalid_data("MOL2 file has no MOLECULE block"))
}

/// Read every molecule in `path` into a Vec.
pub fn read_mol2_all<P: AsRef<Path>>(path: P) -> Result<Vec<Frame>> {
    let file = std::fs::File::open(path.as_ref())?;
    let reader = std::io::BufReader::new(file);
    let mut mr = Mol2Reader::new(reader);
    mr.read_all()
}

/// `FrameReader`-trait wrapper. Each call returns the next molecule or `None`
/// at EOF.
pub struct Mol2Reader<R: BufRead> {
    reader: R,
    pending: Option<String>,
}

impl<R: BufRead> Reader for Mol2Reader<R> {
    type R = R;
    type Frame = Frame;
    fn new(reader: R) -> Self {
        Self {
            reader,
            pending: None,
        }
    }
}

impl<R: BufRead> FrameReader for Mol2Reader<R> {
    fn read_frame(&mut self) -> Result<Option<Frame>> {
        read_one_record(&mut self.reader, &mut self.pending)
    }
}

// ---------------------------------------------------------------------------
// Writer
// ---------------------------------------------------------------------------

/// Write a Frame as a single-molecule `.mol2` file.
pub fn write_mol2<P: AsRef<Path>>(path: P, frame: &Frame) -> Result<()> {
    let file = std::fs::File::create(path.as_ref())?;
    let mut w = BufWriter::new(file);
    write_mol2_frame(&mut w, frame)?;
    w.flush()
}

/// Emit one frame in MOL2 format.
pub fn write_mol2_frame<W: Write>(writer: &mut W, frame: &Frame) -> Result<()> {
    let atoms = frame
        .get("atoms")
        .ok_or_else(|| invalid_data("MOL2 write: frame has no atoms block"))?;
    let n = atoms
        .nrows()
        .ok_or_else(|| invalid_data("MOL2 write: atoms block has no rows"))?;
    let bonds = frame.get("bonds");
    let n_bonds = bonds.and_then(|b| b.nrows()).unwrap_or(0);

    let title = frame
        .meta
        .get("title")
        .cloned()
        .unwrap_or_else(|| "molrs".to_string());

    writeln!(writer, "@<TRIPOS>MOLECULE")?;
    writeln!(writer, "{}", title)?;
    writeln!(writer, "{} {} 0 0 0", n, n_bonds)?;
    writeln!(writer, "SMALL")?;
    let charge_col = atoms.get_float("charge");
    writeln!(
        writer,
        "{}",
        if charge_col.is_some() {
            "USER_CHARGES"
        } else {
            "NO_CHARGES"
        }
    )?;
    writeln!(writer)?;
    writeln!(writer)?;

    writeln!(writer, "@<TRIPOS>ATOM")?;
    let id_col = atoms.get_int("id");
    let name_col = atoms.get_string("name");
    let xs = atoms
        .get_float("x")
        .ok_or_else(|| invalid_data("atoms.x missing"))?;
    let ys = atoms
        .get_float("y")
        .ok_or_else(|| invalid_data("atoms.y missing"))?;
    let zs = atoms
        .get_float("z")
        .ok_or_else(|| invalid_data("atoms.z missing"))?;
    let type_col = atoms.get_string("atom_type");
    let subst_id_col = atoms.get_int("subst_id");
    let subst_name_col = atoms.get_string("subst_name");
    for i in 0..n {
        let id = id_col.map(|c| c[[i]]).unwrap_or((i as I) + 1);
        let name = name_col.map(|c| c[[i]].as_str()).unwrap_or("X");
        let atype = type_col.map(|c| c[[i]].as_str()).unwrap_or("X");
        let sid = subst_id_col.map(|c| c[[i]]).unwrap_or(1);
        let sname = subst_name_col.map(|c| c[[i]].as_str()).unwrap_or("MOL");
        write!(
            writer,
            "{:>7} {:<8} {:>10.4} {:>10.4} {:>10.4} {:<6} {} {:<8}",
            id,
            name,
            xs[[i]],
            ys[[i]],
            zs[[i]],
            atype,
            sid,
            sname
        )?;
        if let Some(c) = charge_col {
            write!(writer, " {:>9.4}", c[[i]])?;
        }
        writeln!(writer)?;
    }

    if let Some(b) = bonds
        && n_bonds > 0
    {
        writeln!(writer, "@<TRIPOS>BOND")?;
        let atomi = b
            .get_uint("atomi")
            .ok_or_else(|| invalid_data("bonds.atomi missing"))?;
        let atomj = b
            .get_uint("atomj")
            .ok_or_else(|| invalid_data("bonds.atomj missing"))?;
        let btype_col = b.get_string("bond_type");
        for i in 0..n_bonds {
            let bt = btype_col.map(|c| c[[i]].as_str()).unwrap_or("1");
            writeln!(
                writer,
                "{:>6} {:>6} {:>6} {}",
                i + 1,
                atomi[[i]] + 1,
                atomj[[i]] + 1,
                bt
            )?;
        }
    }
    Ok(())
}

/// `FrameWriter`-trait wrapper.
pub struct Mol2FrameWriter<W: Write> {
    writer: W,
}

impl<W: Write> Writer for Mol2FrameWriter<W> {
    type W = W;
    type FrameLike = Frame;
    fn new(writer: W) -> Self {
        Self { writer }
    }
}

impl<W: Write> FrameWriter for Mol2FrameWriter<W> {
    fn write_frame(&mut self, frame: &Frame) -> Result<()> {
        write_mol2_frame(&mut self.writer, frame)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    const ETHANE_MIN: &str = "@<TRIPOS>MOLECULE\nETH\n2 1 1 0 0\nSMALL\nNO_CHARGES\n@<TRIPOS>ATOM\n1 C1 0.0 0.0 0.0 c3 1 ETH 0.0\n2 C2 1.5 0.0 0.0 c3 1 ETH 0.0\n@<TRIPOS>BOND\n1 1 2 1\n";

    #[test]
    fn reads_minimal_mol2() {
        let mut reader = Mol2Reader::new(Cursor::new(ETHANE_MIN.as_bytes()));
        let frame = reader.read_frame().unwrap().unwrap();
        let atoms = frame.get("atoms").unwrap();
        assert_eq!(atoms.nrows(), Some(2));
        let xs = atoms.get_float("x").unwrap();
        assert!((xs[[1]] - 1.5).abs() < 1e-9);
        let bonds = frame.get("bonds").unwrap();
        assert_eq!(bonds.nrows(), Some(1));
        let atomi = bonds.get_uint("atomi").unwrap();
        let atomj = bonds.get_uint("atomj").unwrap();
        assert_eq!(atomi[[0]], 0);
        assert_eq!(atomj[[0]], 1);
    }

    #[test]
    fn round_trip_minimal_mol2() {
        let frame = {
            let mut reader = Mol2Reader::new(Cursor::new(ETHANE_MIN.as_bytes()));
            reader.read_frame().unwrap().unwrap()
        };
        let mut buf = Vec::new();
        write_mol2_frame(&mut buf, &frame).unwrap();
        let mut reader2 = Mol2Reader::new(Cursor::new(&buf));
        let frame2 = reader2.read_frame().unwrap().unwrap();
        let xs1 = frame.get("atoms").unwrap().get_float("x").unwrap();
        let xs2 = frame2.get("atoms").unwrap().get_float("x").unwrap();
        for i in 0..xs1.len() {
            assert!((xs1[[i]] - xs2[[i]]).abs() < 1e-3);
        }
    }
}