chematic-mol 0.4.28

//! SDF (Structure-Data File) reader.
//!
//! An SDF file contains one or more MOL V2000 blocks separated by `$$$$`
//! delimiter lines.  Data-field sections between `M  END` and `$$$$` are
//! accepted but ignored.

use chematic_core::Molecule;

use crate::error::MolParseError;
use crate::mol2000::{MolMetadata, parse_mol, parse_mol_with_coords};

/// Iterator over molecules in an SDF string.
///
/// Each call to `next()` returns the next `(Molecule, MolMetadata)` pair
/// parsed from the string, or the first `MolParseError` encountered.
/// Returns `None` when the entire input has been consumed.
pub struct SdfReader<'a> {
    remaining: &'a str,
    current_mol_num: usize,
}

impl<'a> SdfReader<'a> {
    /// Create a new `SdfReader` over the given SDF string.
    pub fn new(input: &'a str) -> Self {
        Self {
            remaining: input,
            current_mol_num: 0,
        }
    }
}

impl<'a> Iterator for SdfReader<'a> {
    type Item = Result<(Molecule, MolMetadata), MolParseError>;

    fn next(&mut self) -> Option<Self::Item> {
        // Skip leading blank lines between records (defensive; well-formed SDF
        // should not have them, but some writers emit a trailing blank).
        while let Some(rest) = self
            .remaining
            .strip_prefix("\r\n")
            .or_else(|| self.remaining.strip_prefix('\n'))
        {
            self.remaining = rest;
        }

        if self.remaining.is_empty() {
            return None;
        }

        self.current_mol_num += 1;

        // Scan line by line so that a `$$$$` substring inside a data value
        // does not trigger a false match.  When the delimiter is found, the
        // mol block runs up to (but excluding) it, and the rest continues
        // after the delimiter line.  When EOF is reached without a delimiter,
        // the entire remainder is treated as a single mol block.
        let mut byte_offset = 0usize;
        let (end_byte, after_delim) = loop {
            let rest = &self.remaining[byte_offset..];
            match rest.find('\n') {
                Some(nl) => {
                    let line = rest[..nl].trim_end_matches('\r');
                    if line == "$$$$" {
                        break (byte_offset, &self.remaining[byte_offset + nl + 1..]);
                    }
                    byte_offset += nl + 1;
                }
                None => {
                    // Last line, no trailing newline.
                    if rest.trim_end_matches('\r') == "$$$$" {
                        break (byte_offset, "");
                    }
                    break (self.remaining.len(), "");
                }
            }
        };

        let mol_block = &self.remaining[..end_byte];
        self.remaining = after_delim;

        if mol_block.trim().is_empty() {
            // Empty block between two `$$$$` lines — skip and try next.
            return self.next();
        }

        Some(parse_mol(mol_block))
    }
}

/// Parse all molecules from an SDF string.
///
/// Stops and returns an error on the first parse failure.
pub fn parse_sdf(input: &str) -> Result<Vec<(Molecule, MolMetadata)>, MolParseError> {
    SdfReader::new(input).collect()
}

// ---------------------------------------------------------------------------
// SdfRecord — molecule + SD data fields
// ---------------------------------------------------------------------------

/// A fully-parsed SDF record: molecule, metadata, 2D coordinates, and SD properties.
///
/// `SdfRecordReader` yields one `SdfRecord` per molecule in an SDF string.
/// The `properties` map corresponds to `> <FieldName>` data blocks.
pub struct SdfRecord {
    /// Parsed molecule (heavy atoms only, no explicit H).
    pub mol: Molecule,
    /// Metadata from the three-line MOL header (name, comment).
    pub meta: MolMetadata,
    /// 2D atom coordinates in Å, indexed by atom position.
    /// Empty when the MOL block contains no coordinate data.
    pub coords: Vec<(f64, f64)>,
    /// SD data fields.  Keys are field names; values are field content.
    /// Multi-line values are joined with `\n`.
    pub properties: std::collections::HashMap<String, String>,
}

/// Iterator over SDF records that also captures SD data fields.
///
/// Unlike [`SdfReader`], this iterator yields [`SdfRecord`] values so that
/// callers can access per-molecule properties (e.g. activity values, MW, etc.).
pub struct SdfRecordReader<'a> {
    remaining: &'a str,
}

impl<'a> SdfRecordReader<'a> {
    /// Create a new reader over the given SDF string.
    pub fn new(input: &'a str) -> Self {
        Self { remaining: input }
    }
}

impl<'a> Iterator for SdfRecordReader<'a> {
    type Item = Result<SdfRecord, MolParseError>;

    fn next(&mut self) -> Option<Self::Item> {
        // Skip leading blank lines.
        while let Some(rest) = self
            .remaining
            .strip_prefix("\r\n")
            .or_else(|| self.remaining.strip_prefix('\n'))
        {
            self.remaining = rest;
        }

        if self.remaining.is_empty() {
            return None;
        }

        // Scan to find the $$$$ delimiter (line-by-line to avoid false matches).
        let mut byte_offset = 0usize;
        let (end_byte, after_delim) = loop {
            let rest = &self.remaining[byte_offset..];
            match rest.find('\n') {
                Some(nl) => {
                    let line = rest[..nl].trim_end_matches('\r');
                    if line == "$$$$" {
                        break (byte_offset, &self.remaining[byte_offset + nl + 1..]);
                    }
                    byte_offset += nl + 1;
                }
                None => {
                    if rest.trim_end_matches('\r') == "$$$$" {
                        break (byte_offset, "");
                    }
                    break (self.remaining.len(), "");
                }
            }
        };

        let block = &self.remaining[..end_byte];
        self.remaining = after_delim;

        if block.trim().is_empty() {
            return self.next();
        }

        // Parse molecule + 2D coordinates.
        let (mol, meta, coords) = match parse_mol_with_coords(block) {
            Ok(triple) => triple,
            Err(e) => return Some(Err(e)),
        };

        // Extract data fields from the part after "M  END".
        let data_part = block
            .find("M  END")
            .map(|pos| &block[pos + 6..])
            .unwrap_or("");
        let properties: std::collections::HashMap<String, String> =
            parse_sd_fields(data_part).into_iter().collect();

        Some(Ok(SdfRecord {
            mol,
            meta,
            coords,
            properties,
        }))
    }
}

/// Parse SD data fields from the section after `M  END`.
///
/// Each field starts with `> <FieldName>` on its own line.  The value is
/// everything on subsequent lines until a blank line (or end of input).
fn parse_sd_fields(data: &str) -> Vec<(String, String)> {
    let mut fields = Vec::new();
    let mut current_key: Option<String> = None;
    let mut current_value_lines: Vec<&str> = Vec::new();

    for raw_line in data.lines() {
        let line = raw_line.trim_end_matches('\r');

        if let Some(key) = parse_sd_field_header(line) {
            // Flush previous field.
            if let Some(k) = current_key.take() {
                fields.push((k, current_value_lines.join("\n")));
                current_value_lines.clear();
            }
            current_key = Some(key);
        } else if line.is_empty() {
            // Blank line ends the current field's value.
            if let Some(k) = current_key.take() {
                fields.push((k, current_value_lines.join("\n")));
                current_value_lines.clear();
            }
        } else if current_key.is_some() {
            current_value_lines.push(line);
        }
    }
    // Flush trailing field with no blank line.
    if let Some(k) = current_key {
        fields.push((k, current_value_lines.join("\n")));
    }

    fields
}

/// Parse `> <FieldName>` header lines, returning the field name or `None`.
fn parse_sd_field_header(line: &str) -> Option<String> {
    // SDF spec: field headers start with "> " and contain the name in `<...>`.
    // We accept "> <Name>" and also ">  <Name>" (extra spaces).
    let rest = line.strip_prefix('>')?;
    let rest = rest.trim();
    let inner = rest.strip_prefix('<')?.strip_suffix('>')?;
    Some(inner.to_string())
}

// ---------------------------------------------------------------------------
// SdfFileReader — streaming reader for file-backed SDF
// ---------------------------------------------------------------------------

/// Streaming SDF iterator over any [`std::io::BufRead`] source.
///
/// Unlike [`SdfRecordReader`] (which requires the entire file in memory as a
/// `&str`), `SdfFileReader` reads one MOL block at a time, suitable for large
/// SDF files.  IO errors are surfaced as [`MolParseError::Io`]; molecule parse
/// errors are returned as `Err` items so the caller can decide to skip or stop.
pub struct SdfFileReader<R: std::io::BufRead> {
    reader: R,
    done: bool,
}

impl<R: std::io::BufRead> SdfFileReader<R> {
    /// Wrap any `BufRead` source (e.g. `BufReader<File>`).
    pub fn new(reader: R) -> Self {
        Self {
            reader,
            done: false,
        }
    }
}

impl<R: std::io::BufRead> Iterator for SdfFileReader<R> {
    type Item = Result<SdfRecord, MolParseError>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.done {
            return None;
        }

        let mut block = String::with_capacity(2048);
        let mut line = String::new();

        loop {
            line.clear();
            match self.reader.read_line(&mut line) {
                Err(e) => {
                    self.done = true;
                    return Some(Err(MolParseError::Io(e.to_string())));
                }
                Ok(0) => {
                    // EOF
                    self.done = true;
                    if block.trim().is_empty() {
                        return None;
                    }
                    // Trailing block without $$$$ delimiter — parse it.
                    break;
                }
                Ok(_) => {
                    let trimmed = line.trim_end_matches(['\r', '\n']);
                    if trimmed == "$$$$" {
                        break;
                    }
                    block.push_str(&line);
                }
            }
        }

        if block.trim().is_empty() {
            // Empty block (e.g. two consecutive $$$$) — advance to next.
            return self.next();
        }

        // Reuse the same parse path as SdfRecordReader.
        let (mol, meta, coords) = match parse_mol_with_coords(&block) {
            Ok(triple) => triple,
            Err(e) => return Some(Err(e)),
        };

        let data_part = block
            .find("M  END")
            .map(|pos| &block[pos + 6..])
            .unwrap_or("");
        let properties: std::collections::HashMap<String, String> =
            parse_sd_fields(data_part).into_iter().collect();

        Some(Ok(SdfRecord {
            mol,
            meta,
            coords,
            properties,
        }))
    }
}

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

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

    const MOL_A: &str = "\
mol_a
  chematic

  2  1  0  0  0  0  0  0  0  0  0 V2000
    0.0000    0.0000    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
    1.0000    0.0000    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
  1  2  1  0
M  END
";

    const MOL_B: &str = "\
mol_b
  chematic

  3  2  0  0  0  0  0  0  0  0  0 V2000
    0.0000    0.0000    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
    1.0000    0.0000    0.0000 N   0  0  0  0  0  0  0  0  0  0  0  0
    2.0000    0.0000    0.0000 O   0  0  0  0  0  0  0  0  0  0  0  0
  1  2  1  0
  2  3  2  0
M  END
";

    fn two_mol_sdf() -> String {
        format!("{MOL_A}$$$$\n{MOL_B}$$$$\n")
    }

    #[test]
    fn test_sdf_reader_two_molecules() {
        let sdf = two_mol_sdf();
        let results: Vec<_> = SdfReader::new(&sdf).collect();
        assert_eq!(results.len(), 2);
        let (mol_a, meta_a) = results[0].as_ref().expect("mol_a parse");
        let (mol_b, meta_b) = results[1].as_ref().expect("mol_b parse");
        assert_eq!(mol_a.atom_count(), 2);
        assert_eq!(mol_a.bond_count(), 1);
        assert_eq!(meta_a.name, "mol_a");
        assert_eq!(mol_b.atom_count(), 3);
        assert_eq!(mol_b.bond_count(), 2);
        assert_eq!(meta_b.name, "mol_b");
    }

    #[test]
    fn test_parse_sdf_all() {
        let sdf = two_mol_sdf();
        let mols = parse_sdf(&sdf).expect("parse_sdf");
        assert_eq!(mols.len(), 2);
    }

    #[test]
    fn test_sdf_reader_single_molecule_no_delimiter() {
        // An SDF with a single molecule that has no trailing $$$$ is still valid.
        let results: Vec<_> = SdfReader::new(MOL_A).collect();
        assert_eq!(results.len(), 1);
        let (mol, _) = results[0].as_ref().expect("parse");
        assert_eq!(mol.atom_count(), 2);
    }

    #[test]
    fn test_sdf_reader_stops_on_error() {
        // Second molecule has a bad counts line; parse_sdf should return Err.
        let bad_sdf = format!("{MOL_A}$$$$\nbad\n  prog\n\n  X  Y\nM  END\n$$$$\n");
        let result = parse_sdf(&bad_sdf);
        assert!(result.is_err());
    }

    #[test]
    fn test_sdf_reader_empty_input() {
        let results: Vec<_> = SdfReader::new("").collect();
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_sdf_reader_names_preserved() {
        let sdf = two_mol_sdf();
        let mols = parse_sdf(&sdf).expect("parse");
        assert_eq!(mols[0].1.name, "mol_a");
        assert_eq!(mols[1].1.name, "mol_b");
    }

    #[test]
    fn test_sdf_with_data_fields() {
        // SDF with data fields between M  END and $$$$ — should be ignored.
        let sdf_with_data = format!("{MOL_A}> <MW>\n44.0\n\n$$$$\n");
        let results: Vec<_> = SdfReader::new(&sdf_with_data).collect();
        assert_eq!(results.len(), 1);
        let (mol, _) = results[0].as_ref().expect("parse");
        assert_eq!(mol.atom_count(), 2);
    }

    #[test]
    fn test_sdf_reader_reports_truncated_large_count_record() {
        let bad_sdf = "\
max_atoms
  chematic

999  0  0  0  0  0  0  0  0  0  0 V2000
$$$$
";
        let mut reader = SdfReader::new(bad_sdf);
        assert!(matches!(
            reader.next(),
            Some(Err(MolParseError::UnexpectedEnd))
        ));
        assert!(reader.next().is_none());
        assert!(matches!(
            parse_sdf(bad_sdf),
            Err(MolParseError::UnexpectedEnd)
        ));
    }

    #[test]
    fn test_sdf_file_reader_streaming() {
        // Use Cursor<Vec<u8>> as an in-memory BufRead to avoid touching the filesystem.
        use std::io::{BufReader, Cursor};
        let sdf = two_mol_sdf();
        let cursor = Cursor::new(sdf.into_bytes());
        let records: Vec<_> = SdfFileReader::new(BufReader::new(cursor))
            .filter_map(|r| r.ok())
            .collect();
        assert_eq!(records.len(), 2);
        assert_eq!(records[0].mol.atom_count(), 2); // mol_a: 2 C atoms
        assert_eq!(records[1].mol.atom_count(), 3); // mol_b: C, N, O
    }

    #[test]
    fn test_sdf_file_reader_skips_empty_block() {
        use std::io::{BufReader, Cursor};
        let sdf = format!("$$$$\n{MOL_A}$$$$\n");
        let cursor = Cursor::new(sdf.into_bytes());
        let records: Vec<_> = SdfFileReader::new(BufReader::new(cursor))
            .filter_map(|r| r.ok())
            .collect();
        assert_eq!(records.len(), 1);
    }

    #[test]
    fn test_sdf_file_reader_malformed_record_yields_err() {
        use std::io::{BufReader, Cursor};

        // 3-record SDF: valid, malformed (bad atom count line), valid
        let malformed = "broken\n  prog\n\n  NOTNUM  0  0 V2000\nM  END\n";
        let sdf = format!("{MOL_A}$$$$\n{malformed}$$$$\n{MOL_B}$$$$\n");
        let cursor = Cursor::new(sdf.into_bytes());
        let results: Vec<_> = SdfFileReader::new(BufReader::new(cursor)).collect();

        // Three items: Ok, Err, Ok
        assert_eq!(results.len(), 3);
        assert!(results[0].is_ok(), "first record should be ok");
        assert!(
            results[1].is_err(),
            "second record should be err (malformed)"
        );
        assert!(results[2].is_ok(), "third record should be ok");
    }
}