Skip to main content

chematic_chem/
standardize.rs

1//! Molecular standardization routines.
2//!
3//! Provides utilities for cleaning up molecular representations:
4//! - Selecting the largest connected fragment.
5//! - Neutralizing simple formal charges.
6
7#![forbid(unsafe_code)]
8
9use std::collections::{HashMap, VecDeque};
10
11use chematic_core::{AtomIdx, BondIdx, Element, Molecule, MoleculeBuilder, validate_valence};
12
13use crate::{hash::mol_hash, hydrogen::remove_hydrogens, tautomer::canonical_tautomer};
14use chematic_smarts::{MatchConfig, find_matches_with_config, parse_smarts};
15
16/// Salt removal catalog: common salt patterns (counterions and solvates).
17///
18/// Each pattern is a (name, SMARTS) tuple for organic and inorganic salts.
19/// Used by [`remove_salts`] to filter out counterions and keep drug-like fragments.
20#[derive(Clone, Debug)]
21pub struct SaltCatalog {
22    /// (name, SMARTS) pairs for salt patterns
23    patterns: Vec<(&'static str, &'static str)>,
24}
25
26impl Default for SaltCatalog {
27    fn default() -> Self {
28        Self::new()
29    }
30}
31
32impl SaltCatalog {
33    /// Create a default salt catalog with common counterions and solvates.
34    pub fn new() -> Self {
35        Self {
36            patterns: vec![
37                // Organic salts (carboxylates, sulfonates, etc.)
38                ("acetate", "[#6](-[#1])(-[#1])-[#6](=[#8])[O-]"),
39                ("formate", "[#6](=[#8])[O-]"),
40                (
41                    "propionate",
42                    "[#6](-[#1])(-[#1])-[#6](-[#1])-[#6](=[#8])[O-]",
43                ),
44                ("benzoate", "c1ccccc1-[#6](=[#8])[O-]"),
45                (
46                    "trifluoroacetate",
47                    "[#9]-[#6](-[#9])(-[#9])-[#6](=[#8])[O-]",
48                ),
49                (
50                    "mesylate",
51                    "[#16](=[#8])(=[#8])-[#8]-[#6](-[#1])(-[#1])-[#1]",
52                ),
53                ("tosylate", "c1ccc(cc1)-[#16](=[#8])(=[#8])-[#8]"),
54                (
55                    "nosylate",
56                    "[#8]-[#6](-[#1])(-[#1])-[#8]-[#16](=[#8])(=[#8])-c1ccc([N+](=O)[O-])cc1",
57                ),
58                ("sulfate", "[#16](=[#8])(=[#8])(-[#8])-[#8]"),
59                ("phosphate", "[#15](=[#8])(-[#8])(-[#8])-[#8]"),
60                (
61                    "citrate",
62                    "[#6](-[#6](=[#8])[O-])(-[#6](=[#8])[O-])-[#6](-[#8])-[#6](=[#8])[O-]",
63                ),
64                (
65                    "tartrate",
66                    "[#6](-[#8])(-[#6](-[#8])-[#6](=[#8])[O-])-[#6](=[#8])[O-]",
67                ),
68                // Inorganic salts (single atoms/small molecules)
69                ("sodium_cation", "[Na+]"),
70                ("potassium_cation", "[K+]"),
71                ("lithium_cation", "[Li+]"),
72                ("calcium_cation", "[Ca+2]"),
73                ("magnesium_cation", "[Mg+2]"),
74                ("chloride_anion", "[Cl-]"),
75                ("bromide_anion", "[Br-]"),
76                ("iodide_anion", "[I-]"),
77                ("fluoride_anion", "[F-]"),
78                ("oxide_anion", "[O-2]"),
79                ("sulfate_anion", "[#16](=[#8])(=[#8])(-[#8])-[#8-]"),
80                ("phosphate_anion", "[#15](=[#8])(-[#8])(-[#8-])-[#8]"),
81                // Solvates and additives
82                ("water", "[#8](-[#1])-[#1]"),
83                (
84                    "dmso",
85                    "[#16](=[#8])(-[#6](-[#1])(-[#1])-[#1])-[#6](-[#1])(-[#1])-[#1]",
86                ),
87                ("methanol", "[#6](-[#1])(-[#1])-[#8]-[#1]"),
88                ("ethanol", "[#6](-[#1])(-[#1])-[#6](-[#1])(-[#1])-[#8]-[#1]"),
89                (
90                    "isopropanol",
91                    "[#6](-[#1])(-[#1])-[#6](-[#8]-[#1])(-[#1])-[#6](-[#1])(-[#1])-[#1]",
92                ),
93                // Rare but important salts
94                ("borate", "[#5](-[#8])(-[#8])-[#8]"),
95                ("ammonium", "[#7+;H0,H1,H2,H3]"),
96            ],
97        }
98    }
99
100    /// Add a custom salt pattern to this catalog.
101    pub fn add(&mut self, name: &'static str, smarts: &'static str) {
102        self.patterns.push((name, smarts));
103    }
104
105    /// Check if a molecule fragment matches any salt pattern.
106    pub fn is_salt(&self, frag: &Molecule) -> bool {
107        // max_matches: Some(1) stops each pattern's VF2 search at the first
108        // embedding instead of enumerating every match — this loop only
109        // needs to know whether a match exists, not how many.
110        let config = MatchConfig {
111            max_visit_budget: Some(1_000_000),
112            max_matches: Some(1),
113            uniquify: false,
114            ..Default::default()
115        };
116        for (_, smarts_str) in &self.patterns {
117            if let Ok(query) = parse_smarts(smarts_str)
118                && !find_matches_with_config(&query, frag, &config).is_empty()
119            {
120                return true;
121            }
122        }
123        false
124    }
125}
126
127/// Find all connected components of `mol` via BFS, sorted descending by size.
128fn connected_components(mol: &Molecule) -> Vec<Vec<AtomIdx>> {
129    let n = mol.atom_count();
130    let mut visited = vec![false; n];
131    let mut components: Vec<Vec<AtomIdx>> = Vec::new();
132
133    for start in 0..n {
134        if visited[start] {
135            continue;
136        }
137        visited[start] = true;
138        let mut component = Vec::new();
139        let mut queue: VecDeque<AtomIdx> = VecDeque::new();
140        queue.push_back(AtomIdx(start as u32));
141
142        while let Some(current) = queue.pop_front() {
143            component.push(current);
144            for (neighbor, _) in mol.neighbors(current) {
145                let ni = neighbor.0 as usize;
146                if !visited[ni] {
147                    visited[ni] = true;
148                    queue.push_back(neighbor);
149                }
150            }
151        }
152        components.push(component);
153    }
154
155    components.sort_by_key(|b| std::cmp::Reverse(b.len()));
156    components
157}
158
159/// Copy bonds from `mol` into `builder` when both endpoints are remapped.
160fn copy_bonds(mol: &Molecule, builder: &mut MoleculeBuilder, remap: &HashMap<AtomIdx, AtomIdx>) {
161    for i in 0..mol.bond_count() {
162        let bond = mol.bond(BondIdx(i as u32));
163        if let (Some(&new_a), Some(&new_b)) = (remap.get(&bond.atom1), remap.get(&bond.atom2)) {
164            let _ = builder.add_bond(new_a, new_b, bond.order);
165        }
166    }
167}
168
169/// Check if a fragment is a common inorganic salt or counterion.
170///
171/// Returns true if fragment matches patterns like: NaCl, KCl, Na+, K+, Cl-, Br-, I-, etc.
172fn is_salt_fragment(frag: &Molecule) -> bool {
173    let n = frag.atom_count();
174
175    // Single atom: check if it's a common counterion
176    if n == 1 {
177        let atom = frag.atom(AtomIdx(0));
178        return matches!(
179            atom.element.atomic_number(),
180            11 | 19 | 37 | 55 |  // Na, K, Rb, Cs (alkali metals)
181            17 | 35 | 53 |       // Cl, Br, I (halogens)
182            8 // O (oxide)
183        );
184    }
185
186    // Two atoms: check for common binary salts (NaCl, KBr, etc.)
187    if n == 2 {
188        let a0 = frag.atom(AtomIdx(0)).element.atomic_number();
189        let a1 = frag.atom(AtomIdx(1)).element.atomic_number();
190        let bond_count = frag.bond_count();
191
192        // Ionic pair (no bond between them) — cation + anion
193        if bond_count == 0 {
194            let metals = [11, 19, 37, 55]; // Na, K, Rb, Cs
195            let nonmetals = [17, 35, 53, 8]; // Cl, Br, I, O
196            return (metals.contains(&a0) && nonmetals.contains(&a1))
197                || (metals.contains(&a1) && nonmetals.contains(&a0));
198        }
199    }
200
201    // Small molecules with only metal/nonmetal atoms (common solvate salts)
202    if n <= 4 {
203        let has_organic = frag.atoms().any(|(_, a)| a.element.atomic_number() == 6);
204        if !has_organic {
205            // Pure inorganic salt (no carbons)
206            return true;
207        }
208    }
209
210    false
211}
212
213/// Return a new `Molecule` with inorganic salts removed, keeping largest organic fragment.
214///
215/// Uses a comprehensive salt catalog (SMARTS patterns) and heuristic detection to identify
216/// and exclude common counterions (Na+, K+, TFA-, acetate, etc.) and inorganic salt fragments.
217///
218/// If no non-salt fragment exists or molecule is empty, returns the original largest fragment.
219pub fn remove_salts(mol: &Molecule) -> Molecule {
220    remove_salts_with_catalog(mol, &SaltCatalog::new())
221}
222
223/// Remove salts using a custom catalog.
224///
225/// # Arguments
226/// - `mol`: the molecule to process
227/// - `catalog`: custom salt catalog (use `SaltCatalog::new()` for default)
228pub fn remove_salts_with_catalog(mol: &Molecule, catalog: &SaltCatalog) -> Molecule {
229    if mol.atom_count() == 0 {
230        return MoleculeBuilder::new().build();
231    }
232
233    let components = connected_components(mol);
234
235    // Find largest non-salt fragment
236    let mut largest_non_salt: Option<&Vec<AtomIdx>> = None;
237    let mut largest_non_salt_size = 0;
238
239    for component in &components {
240        // Extract fragment molecule temporarily to check if it's a salt
241        let mut builder = MoleculeBuilder::new();
242        let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
243        for &old_idx in component {
244            let new_idx = builder.add_atom(mol.atom(old_idx).clone());
245            remap.insert(old_idx, new_idx);
246        }
247        copy_bonds(mol, &mut builder, &remap);
248        let frag = builder.build();
249
250        // Check with catalog first, fall back to heuristic
251        let is_salt = catalog.is_salt(&frag) || is_salt_fragment(&frag);
252
253        // If not a salt and larger than current best, use it
254        if !is_salt && component.len() > largest_non_salt_size {
255            largest_non_salt = Some(component);
256            largest_non_salt_size = component.len();
257        }
258    }
259
260    // Fall back to largest fragment if no non-salt found
261    let component = largest_non_salt.unwrap_or(&components[0]);
262
263    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
264    let mut builder = MoleculeBuilder::new();
265    for &old_idx in component {
266        let new_idx = builder.add_atom(mol.atom(old_idx).clone());
267        remap.insert(old_idx, new_idx);
268    }
269    copy_bonds(mol, &mut builder, &remap);
270    builder.build()
271}
272
273/// Return a new `Molecule` containing only the largest connected fragment.
274///
275/// If the molecule is empty, an empty `Molecule` is returned.
276/// This is an alias for `remove_salts()` for backward compatibility.
277pub fn largest_fragment(mol: &Molecule) -> Molecule {
278    remove_salts(mol)
279}
280
281/// Normalize chemical groups (nitro groups, etc.).
282///
283/// Transforms:
284/// - `[N+](=O)[O-]` → `N(=O)=O` (nitro normalization: N charge 0, O- → double bond)
285///
286/// Returns a new molecule with normalized groups.
287pub fn normalize_groups(mol: &Molecule) -> Molecule {
288    let mut builder = MoleculeBuilder::new();
289    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
290    let mut nitro_atoms = std::collections::HashSet::new();
291    let mut oxide_atoms = std::collections::HashSet::new();
292    let mut azide_atoms = std::collections::HashSet::new();
293    let mut sulfoxide_atoms = std::collections::HashSet::new();
294
295    // First pass: identify functional groups via per-group detectors.
296    for (idx, atom) in mol.atoms() {
297        if atom.element.atomic_number() == 7 && atom.charge == 1 {
298            detect_nitro(mol, idx, atom, &mut nitro_atoms, &mut oxide_atoms);
299            detect_azide(mol, idx, &mut azide_atoms);
300        }
301        if atom.element.atomic_number() == 16 {
302            detect_sulfoxide(mol, idx, &mut sulfoxide_atoms);
303        }
304    }
305
306    // Second pass: copy atoms with normalized charges
307    for (idx, atom) in mol.atoms() {
308        let mut new_atom = atom.clone();
309
310        if nitro_atoms.contains(&idx) {
311            // Neutralize the N and O in nitro group
312            if atom.element.atomic_number() == 7 || atom.element.atomic_number() == 8 {
313                new_atom.charge = 0;
314            }
315        }
316
317        if azide_atoms.contains(&idx) {
318            // Neutralize all N in azide group [N-][N+]#N -> N=N=N
319            if atom.element.atomic_number() == 7 {
320                new_atom.charge = 0;
321            }
322        }
323
324        // Sulfoxide: keep as is (S=O is already correct form)
325
326        let new_idx = builder.add_atom(new_atom);
327        remap.insert(idx, new_idx);
328    }
329
330    // Third pass: copy bonds, normalizing functional groups
331    for i in 0..mol.bond_count() {
332        let bond = mol.bond(chematic_core::BondIdx(i as u32));
333        let mut new_order = bond.order;
334
335        // Nitro groups: convert single N-O (where O is negative) to double
336        if nitro_atoms.contains(&bond.atom1) && nitro_atoms.contains(&bond.atom2) {
337            let a1_is_n = mol.atom(bond.atom1).element.atomic_number() == 7;
338            let a2_is_o = mol.atom(bond.atom2).element.atomic_number() == 8;
339            let a1_is_o = mol.atom(bond.atom1).element.atomic_number() == 8;
340            let a2_is_n = mol.atom(bond.atom2).element.atomic_number() == 7;
341
342            if (a1_is_n
343                && a2_is_o
344                && bond.order == chematic_core::BondOrder::Single
345                && mol.atom(bond.atom2).charge == -1)
346                || (a1_is_o
347                    && a2_is_n
348                    && bond.order == chematic_core::BondOrder::Single
349                    && mol.atom(bond.atom1).charge == -1)
350            {
351                new_order = chematic_core::BondOrder::Double;
352            }
353        }
354
355        // AZIDE normalization: [N-][N+]#N -> N=N=N (convert single to double)
356        if azide_atoms.contains(&bond.atom1) && azide_atoms.contains(&bond.atom2) {
357            let a1_is_n = mol.atom(bond.atom1).element.atomic_number() == 7;
358            let a2_is_n = mol.atom(bond.atom2).element.atomic_number() == 7;
359
360            if a1_is_n && a2_is_n && bond.order == chematic_core::BondOrder::Single {
361                // Convert single bonds in azide to double
362                new_order = chematic_core::BondOrder::Double;
363            }
364        }
365
366        // N-oxide: keep as single bond (already correct after atom charge normalization)
367        if oxide_atoms.contains(&bond.atom1) || oxide_atoms.contains(&bond.atom2) {
368            // No bond order change needed — already single
369        }
370
371        // Sulfoxide: keep as is (S=O is already correct form)
372
373        if let (Some(&new_a1), Some(&new_a2)) = (remap.get(&bond.atom1), remap.get(&bond.atom2)) {
374            let _ = builder.add_bond(new_a1, new_a2, new_order);
375        }
376    }
377
378    builder.build()
379}
380
381// ---------------------------------------------------------------------------
382// Functional group detectors for normalize_groups
383// ---------------------------------------------------------------------------
384
385/// Mark nitro [N+](=O)[O-] and aromatic N-oxide atoms.
386fn detect_nitro(
387    mol: &Molecule,
388    idx: AtomIdx,
389    atom: &chematic_core::Atom,
390    nitro_atoms: &mut std::collections::HashSet<AtomIdx>,
391    oxide_atoms: &mut std::collections::HashSet<AtomIdx>,
392) {
393    let o_nbrs: Vec<_> = mol
394        .neighbors(idx)
395        .filter(|(n, _)| mol.atom(*n).element.atomic_number() == 8)
396        .collect();
397
398    if o_nbrs.len() == 2 {
399        let mut has_double_o = false;
400        let mut has_single_neg_o = false;
401        for (o_idx, bid) in &o_nbrs {
402            let o = mol.atom(*o_idx);
403            let b = mol.bond(*bid);
404            if b.order == chematic_core::BondOrder::Double && o.charge == 0 {
405                has_double_o = true;
406            }
407            if b.order == chematic_core::BondOrder::Single && o.charge == -1 {
408                has_single_neg_o = true;
409                nitro_atoms.insert(*o_idx);
410            }
411        }
412        if has_double_o && has_single_neg_o {
413            nitro_atoms.insert(idx);
414        }
415    } else if let Some((o_idx, bid)) = o_nbrs.first() {
416        let o = mol.atom(*o_idx);
417        let b = mol.bond(*bid);
418        if atom.aromatic && b.order == chematic_core::BondOrder::Single && o.charge == -1 {
419            nitro_atoms.insert(idx);
420            oxide_atoms.insert(*o_idx);
421        }
422    }
423}
424
425/// Mark azide [N-][N+]#N atoms.
426fn detect_azide(
427    mol: &Molecule,
428    idx: AtomIdx,
429    azide_atoms: &mut std::collections::HashSet<AtomIdx>,
430) {
431    let n_nbrs: Vec<_> = mol
432        .neighbors(idx)
433        .filter(|(n, _)| mol.atom(*n).element.atomic_number() == 7)
434        .collect();
435
436    for (n_idx, bid) in &n_nbrs {
437        let n = mol.atom(*n_idx);
438        let b = mol.bond(*bid);
439        if b.order == chematic_core::BondOrder::Triple && n.charge == 0 {
440            for (other_idx, other_bid) in n_nbrs.iter() {
441                if other_idx == n_idx {
442                    continue;
443                }
444                let other = mol.atom(*other_idx);
445                let other_b = mol.bond(*other_bid);
446                if other_b.order == chematic_core::BondOrder::Single && other.charge == -1 {
447                    azide_atoms.insert(idx);
448                    azide_atoms.insert(*n_idx);
449                    azide_atoms.insert(*other_idx);
450                }
451            }
452        }
453    }
454}
455
456/// Mark sulfoxide S=O atoms.
457fn detect_sulfoxide(
458    mol: &Molecule,
459    idx: AtomIdx,
460    sulfoxide_atoms: &mut std::collections::HashSet<AtomIdx>,
461) {
462    for (o_idx, bid) in mol.neighbors(idx) {
463        if mol.atom(o_idx).element.atomic_number() == 8
464            && mol.bond(bid).order == chematic_core::BondOrder::Double
465        {
466            sulfoxide_atoms.insert(idx);
467            sulfoxide_atoms.insert(o_idx);
468        }
469    }
470}
471
472/// Detect if a molecule contains a zwitterion (internal salt).
473///
474/// A zwitterion is defined as having both positive and negative formal charges.
475/// Examples: amino acids in zwitterionic form ([NH3+][COO-]).
476pub fn has_zwitterion(mol: &Molecule) -> bool {
477    let mut has_positive = false;
478    let mut has_negative = false;
479
480    for (_, atom) in mol.atoms() {
481        if atom.charge > 0 {
482            has_positive = true;
483        } else if atom.charge < 0 {
484            has_negative = true;
485        }
486        if has_positive && has_negative {
487            return true;
488        }
489    }
490    false
491}
492
493/// Normalize a zwitterion to neutral form by proton transfer.
494///
495/// For each negatively-charged atom, find the nearest positively-charged atom
496/// and transfer one proton (increase positive charge's H count, decrease negative charge).
497///
498/// Returns a new molecule with normalized charges.
499pub fn normalize_zwitterion(mol: &Molecule) -> Molecule {
500    if !has_zwitterion(mol) {
501        return clone_molecule(mol);
502    }
503
504    let mut modifications: HashMap<AtomIdx, (i8, Option<u8>)> = HashMap::new();
505
506    // Collect positive and negative charge atoms
507    let mut positive_atoms: Vec<AtomIdx> = Vec::new();
508    let mut negative_atoms: Vec<AtomIdx> = Vec::new();
509
510    for i in 0..mol.atom_count() {
511        let idx = AtomIdx(i as u32);
512        let atom = mol.atom(idx);
513        if atom.charge > 0 {
514            positive_atoms.push(idx);
515        } else if atom.charge < 0 {
516            negative_atoms.push(idx);
517        }
518    }
519
520    // For each negative atom, transfer proton from nearest positive atom
521    for &neg_idx in &negative_atoms {
522        if positive_atoms.is_empty() {
523            continue;
524        }
525
526        // Find nearest positive charge (by BFS distance)
527        let mut closest_pos_idx = positive_atoms[0];
528        let mut closest_distance = i32::MAX;
529
530        for &pos_idx in &positive_atoms {
531            if let Some(dist) = bfs_distance(mol, neg_idx, pos_idx)
532                && dist < closest_distance
533            {
534                closest_distance = dist;
535                closest_pos_idx = pos_idx;
536            }
537        }
538
539        // Transfer proton: N+ loses H, O- gains H
540        let neg_atom = mol.atom(neg_idx);
541        let pos_atom = mol.atom(closest_pos_idx);
542
543        // Decrease negative charge
544        let new_neg_charge = neg_atom.charge + 1;
545        let neg_h = neg_atom.hydrogen_count.unwrap_or(0);
546        modifications.insert(neg_idx, (new_neg_charge, Some(neg_h + 1)));
547
548        // Decrease positive charge by decreasing H count
549        let pos_h = pos_atom.hydrogen_count.unwrap_or(0);
550        if pos_h > 0 {
551            let new_pos_charge = pos_atom.charge - 1;
552            modifications.insert(closest_pos_idx, (new_pos_charge, Some(pos_h - 1)));
553        }
554    }
555
556    // Reconstruct molecule with modified charges
557    let mut builder = MoleculeBuilder::new();
558    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
559
560    for i in 0..mol.atom_count() {
561        let old_idx = AtomIdx(i as u32);
562        let mut atom = mol.atom(old_idx).clone();
563        if let Some(&(new_charge, new_h)) = modifications.get(&old_idx) {
564            atom.charge = new_charge;
565            atom.hydrogen_count = new_h;
566        }
567        let new_idx = builder.add_atom(atom);
568        remap.insert(old_idx, new_idx);
569    }
570    copy_bonds(mol, &mut builder, &remap);
571    builder.build()
572}
573
574/// BFS distance between two atoms in a molecule.
575fn bfs_distance(mol: &Molecule, start: AtomIdx, end: AtomIdx) -> Option<i32> {
576    if start == end {
577        return Some(0);
578    }
579
580    let n = mol.atom_count();
581    let mut visited = vec![false; n];
582    let mut queue = std::collections::VecDeque::new();
583    queue.push_back((start, 0));
584    visited[start.0 as usize] = true;
585
586    while let Some((current, dist)) = queue.pop_front() {
587        for (neighbor, _) in mol.neighbors(current) {
588            if neighbor == end {
589                return Some(dist + 1);
590            }
591            let ni = neighbor.0 as usize;
592            if !visited[ni] {
593                visited[ni] = true;
594                queue.push_back((neighbor, dist + 1));
595            }
596        }
597    }
598    None
599}
600
601/// Neutralize simple formal charges in a molecule.
602///
603/// Rules applied:
604/// - `[O-]` with a carbon neighbor → charge 0, +1 H (carboxylate → carboxylic acid).
605/// - `[N+]` with at least one explicit H → charge 0, −1 H (ammonium → amine).
606/// - `[O+]` with at least one explicit H → charge 0, −1 H (protonated ether → ether).
607pub fn neutralize_charges(mol: &Molecule) -> Molecule {
608    let mut modifications: HashMap<AtomIdx, (i8, Option<u8>)> = HashMap::new();
609
610    for i in 0..mol.atom_count() {
611        let idx = AtomIdx(i as u32);
612        let atom = mol.atom(idx);
613        let h = atom.hydrogen_count.unwrap_or(0);
614
615        match (atom.element, atom.charge) {
616            (Element::O, -1) => {
617                let has_c_neighbor = mol
618                    .neighbors(idx)
619                    .any(|(nb, _)| mol.atom(nb).element == Element::C);
620                if has_c_neighbor {
621                    modifications.insert(idx, (0, Some(h + 1)));
622                }
623            }
624            (Element::N, 1) | (Element::O, 1) if h > 0 => {
625                modifications.insert(idx, (0, Some(h - 1)));
626            }
627            _ => {}
628        }
629    }
630
631    let mut builder = MoleculeBuilder::new();
632    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
633    for i in 0..mol.atom_count() {
634        let old_idx = AtomIdx(i as u32);
635        let mut atom = mol.atom(old_idx).clone();
636        if let Some(&(new_charge, new_h)) = modifications.get(&old_idx) {
637            atom.charge = new_charge;
638            atom.hydrogen_count = new_h;
639        }
640        let new_idx = builder.add_atom(atom);
641        remap.insert(old_idx, new_idx);
642    }
643    copy_bonds(mol, &mut builder, &remap);
644    builder.build()
645}
646
647/// Remove all isotope labels from atoms.
648///
649/// Returns a new molecule with `atom.isotope = None` for all atoms.
650pub fn remove_isotopes(mol: &Molecule) -> Molecule {
651    let mut builder = MoleculeBuilder::new();
652    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
653
654    for i in 0..mol.atom_count() {
655        let old_idx = AtomIdx(i as u32);
656        let mut atom = mol.atom(old_idx).clone();
657        atom.isotope = None;
658        let new_idx = builder.add_atom(atom);
659        remap.insert(old_idx, new_idx);
660    }
661    copy_bonds(mol, &mut builder, &remap);
662    builder.build()
663}
664
665/// Remove all stereochemistry from a molecule.
666///
667/// Sets `atom.chirality = Chirality::None` for all atoms and converts
668/// wedge/wedge-hash bonds to single bonds.
669pub fn remove_stereo(mol: &Molecule) -> Molecule {
670    use chematic_core::{BondOrder, Chirality};
671
672    let mut builder = MoleculeBuilder::new();
673    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
674
675    for i in 0..mol.atom_count() {
676        let old_idx = AtomIdx(i as u32);
677        let mut atom = mol.atom(old_idx).clone();
678        atom.chirality = Chirality::None;
679        let new_idx = builder.add_atom(atom);
680        remap.insert(old_idx, new_idx);
681    }
682
683    for i in 0..mol.bond_count() {
684        let bond = mol.bond(BondIdx(i as u32));
685        if let (Some(&new_a), Some(&new_b)) = (remap.get(&bond.atom1), remap.get(&bond.atom2)) {
686            let order = match bond.order {
687                BondOrder::Up | BondOrder::Down => BondOrder::Single,
688                other => other,
689            };
690            let _ = builder.add_bond(new_a, new_b, order);
691        }
692    }
693
694    builder.build()
695}
696
697/// Remove atoms from stereo groups that are no longer chiral, and drop empty groups.
698///
699/// Analog of RDKit PR #9051 ("cleanup of stereogroups and wedges for non-chiral sites").
700/// When molecular operations (fragment removal, chirality clearing, …) leave atoms
701/// without chirality flags, their stereo group membership becomes invalid.  This
702/// function filters each [`StereoGroup`]'s atom list to only those atoms where
703/// `atom.chirality != Chirality::None`, and discards any group that becomes empty.
704pub fn clean_stereo_groups(mol: &Molecule) -> Molecule {
705    use chematic_core::{Chirality, StereoGroup};
706
707    let cleaned: Vec<StereoGroup> = mol
708        .stereo_groups()
709        .iter()
710        .filter_map(|g| {
711            let chiral_atoms: Vec<AtomIdx> = g
712                .atom_indices
713                .iter()
714                .copied()
715                .filter(|&idx| mol.atom(idx).chirality != Chirality::None)
716                .collect();
717            if chiral_atoms.is_empty() {
718                None
719            } else {
720                Some(StereoGroup::new(g.kind.clone(), chiral_atoms))
721            }
722        })
723        .collect();
724
725    let mut out = MoleculeBuilder::from_molecule(mol).build();
726    out.set_stereo_groups(cleaned);
727    out
728}
729
730/// Keep only the largest organic (carbon-containing) fragment.
731///
732/// Removes all inorganic fragments (those without carbon atoms).
733/// Useful for removing metal ions, salts, and other counterions.
734/// Falls back to largest fragment if no organic fragment exists.
735pub fn prefer_organic(mol: &Molecule) -> Molecule {
736    if mol.atom_count() == 0 {
737        return MoleculeBuilder::new().build();
738    }
739
740    let components = connected_components(mol);
741
742    // Find largest organic fragment
743    let mut largest_organic: Option<&Vec<AtomIdx>> = None;
744    let mut largest_organic_size = 0;
745
746    for component in &components {
747        // Check if fragment contains carbon (organic)
748        let has_carbon = component
749            .iter()
750            .any(|&idx| mol.atom(idx).element.atomic_number() == 6);
751
752        if has_carbon && component.len() > largest_organic_size {
753            largest_organic = Some(component);
754            largest_organic_size = component.len();
755        }
756    }
757
758    // Fall back to largest fragment if no organic found
759    let target_component = largest_organic.or_else(|| components.first());
760
761    if let Some(component) = target_component {
762        let mut builder = MoleculeBuilder::new();
763        let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
764        for &old_idx in component {
765            let new_idx = builder.add_atom(mol.atom(old_idx).clone());
766            remap.insert(old_idx, new_idx);
767        }
768        copy_bonds(mol, &mut builder, &remap);
769        builder.build()
770    } else {
771        MoleculeBuilder::new().build()
772    }
773}
774
775/// Reionize a molecule by adjusting protonation to favored forms.
776///
777/// Simple heuristic approach that adjusts charges on common acidic and basic groups:
778/// - Carboxylic acids with OH: deprotonate to COO- (carboxylate)
779/// - Phenols: deprotonate to phenoxide (O-)
780/// - Primary amines: protonate to ammonium (NH3+)
781/// - Imidazoles: protonate to imidazolium (N+)
782///
783/// This is a simplified version suitable for typical organic molecules.
784pub fn reionize(mol: &Molecule) -> Molecule {
785    let mut builder = MoleculeBuilder::new();
786    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
787
788    // Copy all atoms, adjusting charges
789    for i in 0..mol.atom_count() {
790        let idx = AtomIdx(i as u32);
791        let mut atom = mol.atom(idx).clone();
792        let an = atom.element.atomic_number();
793
794        // Check for carboxylic acid or phenol: C=O with O-H or Ar-O-H
795        if an == 8 {
796            // Oxygen: check if it's OH bonded to C
797            if let Some((c_idx, _)) = mol.neighbors(idx).find(|(neighbor, bond_idx)| {
798                mol.bond(*bond_idx).order == chematic_core::BondOrder::Single
799                    && mol.atom(*neighbor).element.atomic_number() == 6
800            }) {
801                // Check if C is aromatic (phenol) or has a double-bonded O (carboxylic acid)
802                let is_aromatic = mol.atom(c_idx).aromatic;
803                let has_double_bonded_o = mol.neighbors(c_idx).any(|(other, bond_idx)| {
804                    mol.bond(bond_idx).order == chematic_core::BondOrder::Double
805                        && mol.atom(other).element.atomic_number() == 8
806                        && other != idx
807                });
808
809                // Only deprotonate if it's a phenol or carboxylic acid, not aliphatic OH
810                if (is_aromatic || has_double_bonded_o) && atom.charge >= 0 {
811                    atom.charge -= 1; // Deprotonate: OH → O-
812                }
813            }
814        }
815
816        // Check for primary/secondary amines (but NOT amides)
817        if an == 7 {
818            // Check if this N is NOT part of an amide (C(=O)-N)
819            let is_amide = mol.neighbors(idx).any(|(neighbor, bond_idx)| {
820                mol.bond(bond_idx).order == chematic_core::BondOrder::Single
821                    && mol.atom(neighbor).element.atomic_number() == 6
822                    && mol.neighbors(neighbor).any(|(o_neighbor, o_bond)| {
823                        mol.bond(o_bond).order == chematic_core::BondOrder::Double
824                            && (mol.atom(o_neighbor).element.atomic_number() == 8
825                                || mol.atom(o_neighbor).element.atomic_number() == 16)
826                    })
827            });
828
829            if !is_amide {
830                let h_count = chematic_core::implicit_hcount(mol, idx);
831                // Protonate free amines only (not amides)
832                if (h_count == 2 || h_count == 1) && atom.charge <= 0 {
833                    atom.charge += 1; // Protonate: NH2 → NH3+
834                }
835            }
836        }
837
838        let new_idx = builder.add_atom(atom);
839        remap.insert(idx, new_idx);
840    }
841
842    copy_bonds(mol, &mut builder, &remap);
843    builder.build()
844}
845
846/// Remove all charges from a molecule by protonation/deprotonation.
847///
848/// Neutralizes positively charged atoms by removing protons and
849/// negatively charged atoms by adding protons. This is an aggressive
850/// neutralization that may create chemically unrealistic structures.
851///
852/// # Note
853/// This differs from [`neutralize_charges`] which uses specific rules.
854/// `uncharge` is a brute-force approach suitable for structure cleanup.
855pub fn uncharge(mol: &Molecule) -> Molecule {
856    let mut builder = MoleculeBuilder::new();
857    let mut remap: HashMap<AtomIdx, AtomIdx> = HashMap::new();
858
859    // Copy all atoms, removing charges
860    for i in 0..mol.atom_count() {
861        let idx = AtomIdx(i as u32);
862        let mut atom = mol.atom(idx).clone();
863        atom.charge = 0; // Force neutral
864        let new_idx = builder.add_atom(atom);
865        remap.insert(idx, new_idx);
866    }
867
868    copy_bonds(mol, &mut builder, &remap);
869    builder.build()
870}
871
872/// One transformation stage in the standardization pipeline.
873#[derive(Clone, Copy, Debug, PartialEq, Eq)]
874#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
875pub enum StandardizationStep {
876    /// Select the largest connected component.
877    LargestFragment,
878    /// Apply simple neutralization rules for common formal charges.
879    NeutralizeCharges,
880    /// Normalize chemical groups (nitro groups, etc.).
881    NormalizeGroups,
882    /// Normalize zwitterionic forms to neutral molecules.
883    ZwitterionNormalization,
884    /// Remove explicit hydrogen atoms.
885    RemoveExplicitHydrogens,
886    /// Canonicalize supported tautomer systems.
887    CanonicalTautomer,
888    /// Keep only the largest non-salt fragment.
889    FragmentParent,
890    /// Neutralize all formal charges.
891    ChargeParent,
892    /// Remove all isotope labels.
893    IsotopeParent,
894    /// Remove all stereochemistry (wedge/wedge-hash bonds, chirality).
895    StereoParent,
896}
897
898impl StandardizationStep {
899    /// Stable machine-readable stage name.
900    pub fn as_str(self) -> &'static str {
901        match self {
902            Self::LargestFragment => "largest_fragment",
903            Self::NeutralizeCharges => "neutralize_charges",
904            Self::NormalizeGroups => "normalize_groups",
905            Self::ZwitterionNormalization => "zwitterion_normalization",
906            Self::RemoveExplicitHydrogens => "remove_explicit_hydrogens",
907            Self::CanonicalTautomer => "canonical_tautomer",
908            Self::FragmentParent => "fragment_parent",
909            Self::ChargeParent => "charge_parent",
910            Self::IsotopeParent => "isotope_parent",
911            Self::StereoParent => "stereo_parent",
912        }
913    }
914}
915
916/// High-level status for a standardization run.
917#[derive(Clone, Copy, Debug, PartialEq, Eq)]
918#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
919pub enum PipelineStatus {
920    /// Pipeline completed and the output is structurally identical to the input.
921    Unchanged,
922    /// Pipeline completed and at least one enabled stage changed the molecule.
923    Modified,
924    /// Pipeline completed, but warnings indicate unsupported or suspicious input features.
925    CompletedWithWarnings,
926}
927
928/// Warning emitted during standardization.
929#[derive(Clone, Debug, PartialEq, Eq)]
930#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
931pub struct StandardizationWarning {
932    /// Stable machine-readable warning code.
933    pub code: String,
934    /// Human-readable detail.
935    pub message: String,
936}
937
938impl StandardizationWarning {
939    fn new(code: &str, message: String) -> Self {
940        Self {
941            code: code.to_string(),
942            message,
943        }
944    }
945}
946
947/// Atom/bond/hash summary before or after a pipeline stage.
948#[derive(Clone, Debug, PartialEq, Eq)]
949#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
950pub struct MoleculeSnapshot {
951    /// Number of atoms in the molecule.
952    pub atoms: usize,
953    /// Number of bonds in the molecule.
954    pub bonds: usize,
955    /// Deterministic structure hash based on canonical SMILES.
956    pub hash: u64,
957}
958
959impl MoleculeSnapshot {
960    fn from_mol(mol: &Molecule) -> Self {
961        Self {
962            atoms: mol.atom_count(),
963            bonds: mol.bond_count(),
964            hash: mol_hash(mol),
965        }
966    }
967}
968
969/// Per-stage audit entry for a standardization run.
970#[derive(Clone, Debug, PartialEq, Eq)]
971#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
972pub struct StandardizationStepReport {
973    /// Pipeline stage.
974    pub step: StandardizationStep,
975    /// Whether the stage was enabled in the config.
976    pub enabled: bool,
977    /// Whether the stage changed the molecule hash.
978    pub changed: bool,
979    /// Molecule summary before the stage.
980    pub before: MoleculeSnapshot,
981    /// Molecule summary after the stage.
982    pub after: MoleculeSnapshot,
983}
984
985/// Full standardization audit result.
986#[derive(Clone, Debug, PartialEq, Eq)]
987#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
988pub struct StandardizationReport {
989    /// Overall status.
990    pub status: PipelineStatus,
991    /// Summary of the input molecule.
992    pub input: MoleculeSnapshot,
993    /// Summary of the output molecule.
994    pub output: MoleculeSnapshot,
995    /// Ordered per-stage results.
996    pub steps: Vec<StandardizationStepReport>,
997    /// Validation and unsupported-feature warnings.
998    pub warnings: Vec<StandardizationWarning>,
999}
1000
1001impl StandardizationReport {
1002    /// Returns `true` if the molecule changed at any enabled stage.
1003    pub fn changed(&self) -> bool {
1004        self.input.hash != self.output.hash
1005    }
1006}
1007
1008/// Handling strategy for zwitterions (internal salts).
1009#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
1010pub enum ZwitterionHandling {
1011    /// Keep zwitterionic form as-is.
1012    Keep,
1013    /// Normalize to neutral form via proton transfer.
1014    #[default]
1015    Normalize,
1016}
1017
1018/// Options for molecular standardization.
1019///
1020/// Controls which cleaning transformations are applied in a standardization pipeline.
1021#[derive(Clone, Debug)]
1022pub struct StandardizeOptions {
1023    /// Convert to canonical tautomer. Default: `true`.
1024    pub canonical_tautomer: bool,
1025    /// Neutralize simple formal charges. Default: `true`.
1026    pub neutralize_charges: bool,
1027    /// Remove explicit hydrogen atoms. Default: `true`.
1028    pub remove_explicit_h: bool,
1029    /// Keep only the largest connected fragment. Default: `false`.
1030    pub largest_fragment_only: bool,
1031    /// Handle zwitterions (internal salts). Default: `Normalize`.
1032    pub zwitterion_handling: ZwitterionHandling,
1033}
1034
1035impl Default for StandardizeOptions {
1036    fn default() -> Self {
1037        Self {
1038            canonical_tautomer: true,
1039            neutralize_charges: true,
1040            remove_explicit_h: true,
1041            largest_fragment_only: false,
1042            zwitterion_handling: ZwitterionHandling::Normalize,
1043        }
1044    }
1045}
1046
1047/// RDKit-style standardization pipeline with an auditable report.
1048#[derive(Clone, Debug, Default)]
1049pub struct StandardizationPipeline {
1050    options: StandardizeOptions,
1051}
1052
1053impl StandardizationPipeline {
1054    /// Create a pipeline from explicit options.
1055    pub fn new(options: StandardizeOptions) -> Self {
1056        Self { options }
1057    }
1058
1059    /// Borrow the pipeline options.
1060    pub fn options(&self) -> &StandardizeOptions {
1061        &self.options
1062    }
1063
1064    /// Standardize a molecule and return both the output molecule and an audit report.
1065    pub fn run(&self, mol: &Molecule) -> (Molecule, StandardizationReport) {
1066        let input = MoleculeSnapshot::from_mol(mol);
1067        let mut current = clone_molecule(mol);
1068        let mut steps = Vec::new();
1069        let mut warnings = detect_initial_warnings(mol);
1070
1071        // Disconnect metals early (remove dative/coordinate bonds)
1072        let has_metals = current.atoms().any(|(_, a)| is_metal(a.element));
1073        if has_metals {
1074            current = disconnect_metals(&current);
1075        }
1076
1077        // Apply NeutralizeCharges BEFORE LargestFragment to ensure predictable fragment selection.
1078        // Example: [NH3+].[Cl-] should be neutralized first to [NH3].[Cl-], then largest fragment.
1079        current = self.apply_stage(
1080            current,
1081            StandardizationStep::NeutralizeCharges,
1082            self.options.neutralize_charges,
1083            neutralize_charges,
1084            &mut steps,
1085            &mut warnings,
1086        );
1087        current = self.apply_stage(
1088            current,
1089            StandardizationStep::LargestFragment,
1090            self.options.largest_fragment_only,
1091            largest_fragment,
1092            &mut steps,
1093            &mut warnings,
1094        );
1095        let zwitterion_enabled = self.options.zwitterion_handling == ZwitterionHandling::Normalize;
1096        current = self.apply_stage(
1097            current,
1098            StandardizationStep::ZwitterionNormalization,
1099            zwitterion_enabled,
1100            normalize_zwitterion,
1101            &mut steps,
1102            &mut warnings,
1103        );
1104        current = self.apply_stage(
1105            current,
1106            StandardizationStep::RemoveExplicitHydrogens,
1107            self.options.remove_explicit_h,
1108            remove_hydrogens,
1109            &mut steps,
1110            &mut warnings,
1111        );
1112        current = self.apply_stage(
1113            current,
1114            StandardizationStep::CanonicalTautomer,
1115            self.options.canonical_tautomer,
1116            canonical_tautomer,
1117            &mut steps,
1118            &mut warnings,
1119        );
1120
1121        let output = MoleculeSnapshot::from_mol(&current);
1122        // Status depends only on structure change (hash), not on warnings.
1123        // Warnings are reported separately for the user to inspect.
1124        let status = if input.hash == output.hash {
1125            PipelineStatus::Unchanged
1126        } else if !warnings.is_empty() {
1127            PipelineStatus::CompletedWithWarnings
1128        } else {
1129            PipelineStatus::Modified
1130        };
1131
1132        (
1133            current,
1134            StandardizationReport {
1135                status,
1136                input,
1137                output,
1138                steps,
1139                warnings,
1140            },
1141        )
1142    }
1143
1144    fn apply_stage(
1145        &self,
1146        current: Molecule,
1147        step: StandardizationStep,
1148        enabled: bool,
1149        f: fn(&Molecule) -> Molecule,
1150        steps: &mut Vec<StandardizationStepReport>,
1151        warnings: &mut Vec<StandardizationWarning>,
1152    ) -> Molecule {
1153        let before = MoleculeSnapshot::from_mol(&current);
1154        let next = if enabled {
1155            f(&current)
1156        } else {
1157            clone_molecule(&current)
1158        };
1159        let after = MoleculeSnapshot::from_mol(&next);
1160        steps.push(StandardizationStepReport {
1161            step,
1162            enabled,
1163            changed: before.hash != after.hash,
1164            before,
1165            after,
1166        });
1167        if enabled {
1168            append_valence_warnings(step, &next, warnings);
1169        }
1170        next
1171    }
1172}
1173
1174fn clone_molecule(mol: &Molecule) -> Molecule {
1175    MoleculeBuilder::from_molecule(mol).build()
1176}
1177
1178fn detect_initial_warnings(mol: &Molecule) -> Vec<StandardizationWarning> {
1179    let mut warnings = Vec::new();
1180    // Metal disconnection is now handled in the pipeline, so we don't warn about it
1181    let valence_errors = validate_valence(mol);
1182    if !valence_errors.is_empty() {
1183        warnings.push(StandardizationWarning::new(
1184            "input_valence_validation_failed",
1185            format!(
1186                "input molecule has {} valence validation issue(s)",
1187                valence_errors.len()
1188            ),
1189        ));
1190    }
1191    warnings
1192}
1193
1194fn append_valence_warnings(
1195    step: StandardizationStep,
1196    mol: &Molecule,
1197    warnings: &mut Vec<StandardizationWarning>,
1198) {
1199    let errors = validate_valence(mol);
1200    if errors.is_empty() {
1201        return;
1202    }
1203    warnings.push(StandardizationWarning::new(
1204        "valence_validation_failed",
1205        format!(
1206            "{} produced {} valence validation issue(s)",
1207            step.as_str(),
1208            errors.len()
1209        ),
1210    ));
1211}
1212
1213/// Disconnect metal-nonmetal bonds by removing dative/coordinate bonds to metals.
1214///
1215/// Iterates through all atoms; if a metal is found, removes all bonds between
1216/// that metal and organic/inorganic atoms. Returns the molecule with metal
1217/// coordination bonds severed.
1218fn disconnect_metals(mol: &Molecule) -> Molecule {
1219    let mut builder = MoleculeBuilder::new();
1220
1221    // Copy all atoms
1222    for i in 0..mol.atom_count() {
1223        builder.add_atom(mol.atom(AtomIdx(i as u32)).clone());
1224    }
1225
1226    // Copy only non-metal bonds
1227    for i in 0..mol.bond_count() {
1228        let bond = mol.bond(BondIdx(i as u32));
1229        let atom1_is_metal = is_metal(mol.atom(bond.atom1).element);
1230        let atom2_is_metal = is_metal(mol.atom(bond.atom2).element);
1231
1232        // Skip bonds where at least one atom is a metal
1233        if !atom1_is_metal && !atom2_is_metal {
1234            builder.add_bond(bond.atom1, bond.atom2, bond.order).ok();
1235        }
1236    }
1237
1238    builder.build()
1239}
1240
1241fn is_metal(element: Element) -> bool {
1242    matches!(
1243        element.atomic_number(),
1244        3 | 4
1245            | 11 | 12 | 13
1246            | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31
1247            | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50
1248            | 55 | 56 | 57..=71
1249            | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83
1250            | 87 | 88 | 89..=103
1251            | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | 113 | 114 | 115 | 116
1252    )
1253}
1254
1255/// Apply a series of standardization steps to a molecule.
1256///
1257/// Transformations are applied in this order:
1258/// 1. If `largest_fragment_only`, select the largest connected component.
1259/// 2. If `neutralize_charges`, neutralize simple charges.
1260/// 3. If `remove_explicit_h`, remove explicit H atoms.
1261/// 4. If `canonical_tautomer`, convert to the canonical tautomer.
1262///
1263/// Useful for cleaning pasted structures or database entries.
1264pub fn standardize(mol: &Molecule, opts: &StandardizeOptions) -> Molecule {
1265    StandardizationPipeline::new(opts.clone()).run(mol).0
1266}
1267
1268#[cfg(test)]
1269mod tests {
1270    use super::*;
1271    use chematic_smiles::parse;
1272
1273    #[test]
1274    fn largest_fragment_two_fragments_picks_larger() {
1275        // "CC.CCC" — ethane (2 C) and propane (3 C)
1276        let mol = parse("CC.CCC").unwrap();
1277        let result = largest_fragment(&mol);
1278        assert_eq!(result.atom_count(), 3, "should keep propane (3 C)");
1279    }
1280
1281    #[test]
1282    fn largest_fragment_single_fragment_unchanged() {
1283        // "CC" — ethane, only one fragment
1284        let mol = parse("CC").unwrap();
1285        let result = largest_fragment(&mol);
1286        assert_eq!(result.atom_count(), 2);
1287    }
1288
1289    #[test]
1290    fn largest_fragment_keeps_benzene_over_ethane() {
1291        // "CC.c1ccccc1" — ethane (2 C) vs benzene (6 C)
1292        let mol = parse("CC.c1ccccc1").unwrap();
1293        let result = largest_fragment(&mol);
1294        assert_eq!(result.atom_count(), 6, "should keep benzene (6 atoms)");
1295    }
1296
1297    #[test]
1298    fn largest_fragment_ionic_pair_keeps_one_atom() {
1299        // "[Na+].[Cl-]" — both fragments are single atoms; either is fine
1300        let mol = parse("[Na+].[Cl-]").unwrap();
1301        let result = largest_fragment(&mol);
1302        assert_eq!(result.atom_count(), 1);
1303    }
1304
1305    #[test]
1306    fn neutralize_neutral_molecule_unchanged() {
1307        // "CC" is already neutral; no atom should gain/lose charge
1308        let mol = parse("CC").unwrap();
1309        let result = neutralize_charges(&mol);
1310        for i in 0..result.atom_count() {
1311            let atom = result.atom(AtomIdx(i as u32));
1312            assert_eq!(atom.charge, 0, "all atoms should remain neutral");
1313        }
1314    }
1315
1316    #[test]
1317    fn neutralize_acetate_oxygen() {
1318        // "CC(=O)[O-]" — acetate; the [O-] should become neutral with H added
1319        let mol = parse("CC(=O)[O-]").unwrap();
1320        let result = neutralize_charges(&mol);
1321
1322        // Find the oxygen that was originally [O-]: it should now have charge 0
1323        // and hydrogen_count == Some(1).
1324        let neutralized_o = (0..result.atom_count())
1325            .map(|i| result.atom(AtomIdx(i as u32)))
1326            .find(|a| a.element == Element::O && a.hydrogen_count == Some(1));
1327
1328        assert!(
1329            neutralized_o.is_some(),
1330            "neutralized [O-] should have hydrogen_count == Some(1)"
1331        );
1332        assert_eq!(
1333            neutralized_o.unwrap().charge,
1334            0,
1335            "neutralized [O-] should have charge == 0"
1336        );
1337    }
1338
1339    #[test]
1340    fn standardize_with_defaults() {
1341        // "CC(=O)[O-]" — acetate ion
1342        let mol = parse("CC(=O)[O-]").unwrap();
1343        let opts = StandardizeOptions::default();
1344        let result = standardize(&mol, &opts);
1345
1346        // With default options, remove_explicit_h will be applied.
1347        // Check that [O-] was neutralized (charge should be 0).
1348        let has_neutral_o = (0..result.atom_count())
1349            .map(|i| result.atom(AtomIdx(i as u32)))
1350            .any(|a| a.element == Element::O && a.charge == 0);
1351        assert!(has_neutral_o, "acetate oxygen should be neutralized");
1352
1353        // Should have at least 3 atoms (C, C, O) with no explicit H
1354        assert!(
1355            result.atom_count() >= 3,
1356            "should have at least 3 atoms after standardization"
1357        );
1358    }
1359
1360    #[test]
1361    fn standardize_skip_largest_fragment() {
1362        // "CC.CCC" — ethane and propane
1363        let mol = parse("CC.CCC").unwrap();
1364        let opts = StandardizeOptions {
1365            zwitterion_handling: ZwitterionHandling::Normalize,
1366            largest_fragment_only: false,
1367            ..Default::default()
1368        };
1369        let result = standardize(&mol, &opts);
1370
1371        // Should keep both fragments
1372        assert_eq!(
1373            result.atom_count(),
1374            5,
1375            "should keep both fragments when largest_fragment_only=false"
1376        );
1377    }
1378
1379    #[test]
1380    fn pipeline_report_tracks_enabled_stage_changes() {
1381        let mol = parse("CC.CCC").unwrap();
1382        let pipeline = StandardizationPipeline::new(StandardizeOptions {
1383            largest_fragment_only: true,
1384            neutralize_charges: false,
1385            remove_explicit_h: false,
1386            canonical_tautomer: false,
1387            zwitterion_handling: ZwitterionHandling::Keep,
1388        });
1389
1390        let (result, report) = pipeline.run(&mol);
1391
1392        assert_eq!(result.atom_count(), 3);
1393        assert_eq!(report.status, PipelineStatus::Modified);
1394        assert!(report.changed());
1395        assert_eq!(report.steps.len(), 5);
1396        // NeutralizeCharges is applied first (not enabled, so no change)
1397        assert_eq!(report.steps[0].step, StandardizationStep::NeutralizeCharges);
1398        assert!(!report.steps[0].enabled);
1399        // LargestFragment is applied second and is enabled
1400        assert_eq!(report.steps[1].step, StandardizationStep::LargestFragment);
1401        assert!(report.steps[1].enabled);
1402        assert!(report.steps[1].changed);
1403    }
1404
1405    #[test]
1406    fn pipeline_report_marks_unchanged_clean_molecule() {
1407        let mol = parse("CC").unwrap();
1408        let pipeline = StandardizationPipeline::new(StandardizeOptions {
1409            canonical_tautomer: false,
1410            neutralize_charges: false,
1411            remove_explicit_h: false,
1412            largest_fragment_only: false,
1413            zwitterion_handling: ZwitterionHandling::Keep,
1414        });
1415
1416        let (_result, report) = pipeline.run(&mol);
1417
1418        assert_eq!(report.status, PipelineStatus::Unchanged);
1419        assert!(!report.changed());
1420        assert!(report.warnings.is_empty());
1421        assert!(report.steps.iter().all(|s| !s.enabled && !s.changed));
1422    }
1423
1424    #[test]
1425    fn pipeline_report_disconnects_metal_bonds() {
1426        let mol = parse("[Na]OC").unwrap();
1427        assert_eq!(mol.bond_count(), 2, "input has Na-O and O-C bonds");
1428
1429        let pipeline = StandardizationPipeline::new(StandardizeOptions {
1430            canonical_tautomer: false,
1431            neutralize_charges: false,
1432            remove_explicit_h: false,
1433            largest_fragment_only: false,
1434            zwitterion_handling: ZwitterionHandling::Keep,
1435        });
1436
1437        let (result, _report) = pipeline.run(&mol);
1438
1439        // Metal disconnection should run automatically, removing the Na-O bond
1440        assert_eq!(result.bond_count(), 1, "Na-O bond should be disconnected");
1441        // The remaining bond should be O-C
1442        assert!(
1443            result.bond(BondIdx(0)).atom1.0 < 3 && result.bond(BondIdx(0)).atom2.0 < 3,
1444            "remaining bond should connect organic atoms"
1445        );
1446    }
1447
1448    #[test]
1449    fn bug3_ionic_pair_neutralize_before_largest_fragment() {
1450        // BUG #3 Fix Verification: NeutralizeCharges must run BEFORE LargestFragment
1451        // Example: [NH3+].[OH-] (ammonium hydroxide)
1452        // NH3+ will be neutralized (reduce H by 1)
1453        // After neutralization: [NH2+].[OH-] - still two fragments
1454        // LargestFragment then picks the larger one
1455        // Key: stage order affects which fragment is selected
1456        let mol = parse("[NH3+].[OH-]").unwrap();
1457        let pipeline = StandardizationPipeline::new(StandardizeOptions {
1458            largest_fragment_only: true,
1459            neutralize_charges: true,
1460            remove_explicit_h: false,
1461            canonical_tautomer: false,
1462            zwitterion_handling: ZwitterionHandling::Normalize,
1463        });
1464
1465        let (_result, report) = pipeline.run(&mol);
1466
1467        // Verify step order: NeutralizeCharges MUST come before LargestFragment
1468        assert_eq!(report.steps.len(), 5, "Should have 5 steps in pipeline");
1469        assert_eq!(
1470            report.steps[0].step,
1471            StandardizationStep::NeutralizeCharges,
1472            "NeutralizeCharges must be step 0"
1473        );
1474        assert_eq!(
1475            report.steps[1].step,
1476            StandardizationStep::LargestFragment,
1477            "LargestFragment must be step 1"
1478        );
1479
1480        // The test passes if step order is correct and the pipeline runs without error
1481        assert!(report.changed(), "Pipeline should report changes");
1482        assert_eq!(
1483            report.status,
1484            PipelineStatus::Modified,
1485            "Should be marked as Modified"
1486        );
1487    }
1488
1489    // ── Parent structure extraction tests ────────────────────────────────────
1490
1491    #[test]
1492    fn remove_isotopes_strips_isotope_labels() {
1493        // "[13C]CC" has one 13C isotope label
1494        let mol = parse("[13C]CC").unwrap();
1495        let result = remove_isotopes(&mol);
1496        for i in 0..result.atom_count() {
1497            assert_eq!(
1498                result.atom(chematic_core::AtomIdx(i as u32)).isotope,
1499                None,
1500                "atom {} should have no isotope",
1501                i
1502            );
1503        }
1504    }
1505
1506    #[test]
1507    fn remove_isotopes_preserves_structure() {
1508        // "[13C]CC" and "CC" should be structurally identical after isotope removal
1509        let mol = parse("[13C]CC").unwrap();
1510        let result = remove_isotopes(&mol);
1511        assert_eq!(result.atom_count(), 3, "atom count preserved");
1512        assert_eq!(result.bond_count(), 2, "bond count preserved");
1513    }
1514
1515    #[test]
1516    fn remove_stereo_strips_chirality() {
1517        // "N[C@@H](C)C(=O)O" — alanine with (S) stereochemistry
1518        let mol = parse("N[C@@H](C)C(=O)O").unwrap();
1519        let result = remove_stereo(&mol);
1520        for i in 0..result.atom_count() {
1521            use chematic_core::Chirality;
1522            assert_eq!(
1523                result.atom(chematic_core::AtomIdx(i as u32)).chirality,
1524                Chirality::None,
1525                "atom {} should have no chirality",
1526                i
1527            );
1528        }
1529    }
1530
1531    #[test]
1532    fn remove_stereo_converts_wedge_bonds_to_single() {
1533        // "C[C@H](O)C" — methylcarbinol with wedge stereochemistry
1534        // After remove_stereo, Up/Down bonds should become Single
1535        let mol = parse("C[C@H](O)C").unwrap();
1536        let result = remove_stereo(&mol);
1537        for i in 0..result.bond_count() {
1538            use chematic_core::BondOrder;
1539            let bond = result.bond(chematic_core::BondIdx(i as u32));
1540            assert_ne!(
1541                bond.order,
1542                BondOrder::Up,
1543                "bond {} should not be Up after stereo removal",
1544                i
1545            );
1546            assert_ne!(
1547                bond.order,
1548                BondOrder::Down,
1549                "bond {} should not be Down after stereo removal",
1550                i
1551            );
1552        }
1553    }
1554
1555    #[test]
1556    fn remove_stereo_preserves_structure() {
1557        // "N[C@@H](C)C(=O)O" (alanine) should keep 6 atoms, 5 bonds after stereo removal
1558        let mol = parse("N[C@@H](C)C(=O)O").unwrap();
1559        let result = remove_stereo(&mol);
1560        assert_eq!(result.atom_count(), 6, "atom count preserved");
1561        assert_eq!(result.bond_count(), 5, "bond count preserved");
1562    }
1563
1564    #[test]
1565    fn parent_variant_step_names_distinct() {
1566        // Verify that the 4 new parent variants have distinct step names
1567        let frag_parent = StandardizationStep::FragmentParent;
1568        let charge_parent = StandardizationStep::ChargeParent;
1569        let isotope_parent = StandardizationStep::IsotopeParent;
1570        let stereo_parent = StandardizationStep::StereoParent;
1571
1572        assert_eq!(frag_parent.as_str(), "fragment_parent");
1573        assert_eq!(charge_parent.as_str(), "charge_parent");
1574        assert_eq!(isotope_parent.as_str(), "isotope_parent");
1575        assert_eq!(stereo_parent.as_str(), "stereo_parent");
1576    }
1577
1578    // ── Cleanup transform tests (B3) ────────────────────────────────────────
1579
1580    #[test]
1581    fn prefer_organic_removes_inorganic_salts() {
1582        // "CCO.[Na+].[Cl-]" — ethanol + sodium chloride
1583        let mol = parse("CCO.[Na+].[Cl-]").unwrap();
1584        assert_eq!(mol.atom_count(), 5, "input has CCO + Na + Cl");
1585
1586        let result = prefer_organic(&mol);
1587
1588        // Should keep only the organic ethanol fragment
1589        assert_eq!(result.atom_count(), 3, "should keep only ethanol (C, C, O)");
1590    }
1591
1592    #[test]
1593    fn prefer_organic_keeps_organic_if_no_inorganic() {
1594        // "CC" — ethane only
1595        let mol = parse("CC").unwrap();
1596        let result = prefer_organic(&mol);
1597        assert_eq!(result.atom_count(), 2, "ethane unchanged");
1598    }
1599
1600    #[test]
1601    fn prefer_organic_falls_back_to_largest() {
1602        // "C.C.C" — three separate carbons (all organic)
1603        let mol = parse("C.C.C").unwrap();
1604        let result = prefer_organic(&mol);
1605        // Should keep the largest organic fragment (any one of them, but all are size 1)
1606        assert_eq!(
1607            result.atom_count(),
1608            1,
1609            "falls back to largest fragment (one C)"
1610        );
1611    }
1612
1613    #[test]
1614    fn uncharge_neutralizes_all_charges() {
1615        // "[NH4+].[OH-]" — ammonium hydroxide
1616        let mol = parse("[NH4+].[OH-]").unwrap();
1617        assert!(
1618            mol.atoms().any(|(_, a)| a.charge != 0),
1619            "input has charged atoms"
1620        );
1621
1622        let result = uncharge(&mol);
1623
1624        // All atoms should be neutral
1625        for (_, atom) in result.atoms() {
1626            assert_eq!(atom.charge, 0, "all atoms should be neutral");
1627        }
1628    }
1629
1630    #[test]
1631    fn reionize_deprotonates_carboxylic_acids() {
1632        // "CC(=O)O" — acetic acid
1633        let mol = parse("CC(=O)O").unwrap();
1634
1635        let result = reionize(&mol);
1636
1637        // Should have a negatively charged oxygen (carboxylate anion)
1638        let has_negative_oxygen = result
1639            .atoms()
1640            .any(|(_, a)| a.element.atomic_number() == 8 && a.charge < 0);
1641
1642        assert!(
1643            has_negative_oxygen,
1644            "reionize should deprotonate carboxylic acids"
1645        );
1646    }
1647
1648    #[test]
1649    fn reionize_protonates_amines() {
1650        // "CC(N)C" — secondary amine
1651        let mol = parse("CC(N)C").unwrap();
1652
1653        let result = reionize(&mol);
1654
1655        // Should have a positively charged nitrogen (ammonium)
1656        let has_positive_nitrogen = result
1657            .atoms()
1658            .any(|(_, a)| a.element.atomic_number() == 7 && a.charge > 0);
1659
1660        assert!(has_positive_nitrogen, "reionize should protonate amines");
1661    }
1662
1663    #[test]
1664    fn reionize_protects_amide_nitrogen() {
1665        // BUG FIX #2: Amide nitrogen should NOT be protonated
1666        // "CC(=O)N" — primary amide
1667        let mol = parse("CC(=O)N").unwrap();
1668
1669        let result = reionize(&mol);
1670
1671        // Should NOT have a positively charged nitrogen
1672        let has_positive_nitrogen = result
1673            .atoms()
1674            .any(|(_, a)| a.element.atomic_number() == 7 && a.charge > 0);
1675
1676        assert!(
1677            !has_positive_nitrogen,
1678            "reionize should NOT protonate amide nitrogen"
1679        );
1680    }
1681
1682    #[test]
1683    fn reionize_protects_thioamide_nitrogen() {
1684        // BUG FIX #2: Thioamide nitrogen should also NOT be protonated
1685        // "CC(=S)N" — thioamide
1686        let mol = parse("CC(=S)N").unwrap();
1687
1688        let result = reionize(&mol);
1689
1690        // Should NOT have a positively charged nitrogen
1691        let has_positive_nitrogen = result
1692            .atoms()
1693            .any(|(_, a)| a.element.atomic_number() == 7 && a.charge > 0);
1694
1695        assert!(
1696            !has_positive_nitrogen,
1697            "reionize should NOT protonate thioamide nitrogen (C=S conjugation)"
1698        );
1699    }
1700
1701    // B2: normalize_groups expansion tests
1702
1703    #[test]
1704    fn normalize_groups_nitro() {
1705        // Standard nitro group: [N+](=O)[O-]
1706        let mol = parse("C[N+](=O)[O-]").unwrap();
1707        let result = normalize_groups(&mol);
1708
1709        // All atoms should be neutral after normalization
1710        let all_neutral = result.atoms().all(|(_, a)| a.charge == 0);
1711        assert!(all_neutral, "nitro group should be neutralized");
1712
1713        // N-O bonds should be double
1714        let mut has_double_bond = false;
1715        for (_, bond) in result.bonds() {
1716            let a1 = result.atom(bond.atom1);
1717            let a2 = result.atom(bond.atom2);
1718            if ((a1.element.atomic_number() == 7 && a2.element.atomic_number() == 8)
1719                || (a1.element.atomic_number() == 8 && a2.element.atomic_number() == 7))
1720                && bond.order == chematic_core::BondOrder::Double
1721            {
1722                has_double_bond = true;
1723            }
1724        }
1725        assert!(has_double_bond, "nitro should have N=O double bond");
1726    }
1727
1728    #[test]
1729    fn normalize_groups_azide() {
1730        // Azide: [N-][N+]#N
1731        let mol = parse("[N-][N+]#N").unwrap();
1732        let result = normalize_groups(&mol);
1733
1734        // All atoms should be neutral
1735        let all_neutral = result.atoms().all(|(_, a)| a.charge == 0);
1736        assert!(all_neutral, "azide should be neutralized");
1737
1738        // Check for N=N bonds (converted from single)
1739        let mut has_double_bond_count = 0;
1740        for (_, bond) in result.bonds() {
1741            let a1 = result.atom(bond.atom1);
1742            let a2 = result.atom(bond.atom2);
1743            if a1.element.atomic_number() == 7
1744                && a2.element.atomic_number() == 7
1745                && bond.order == chematic_core::BondOrder::Double
1746            {
1747                has_double_bond_count += 1;
1748            }
1749        }
1750        assert!(
1751            has_double_bond_count > 0,
1752            "azide should have N=N double bonds after normalization"
1753        );
1754    }
1755
1756    #[test]
1757    fn normalize_groups_sulfoxide() {
1758        // Sulfoxide: S(=O)(C)(C)
1759        let mol = parse("C[S](=O)C").unwrap();
1760        let result = normalize_groups(&mol);
1761
1762        // Sulfoxide structure should remain (S=O is already correct form)
1763        let mut has_s_double_o = false;
1764        for (_, bond) in result.bonds() {
1765            let a1 = result.atom(bond.atom1);
1766            let a2 = result.atom(bond.atom2);
1767            if ((a1.element.atomic_number() == 16 && a2.element.atomic_number() == 8)
1768                || (a1.element.atomic_number() == 8 && a2.element.atomic_number() == 16))
1769                && bond.order == chematic_core::BondOrder::Double
1770            {
1771                has_s_double_o = true;
1772            }
1773        }
1774        assert!(has_s_double_o, "sulfoxide should have S=O double bond");
1775    }
1776
1777    // ── RDKit PR #9051: StereoGroup cleanup for non-chiral atoms ────────────
1778
1779    #[test]
1780    fn remove_stereo_clears_stereo_groups() {
1781        // remove_stereo must produce a molecule with no stereo groups (RDKit PR #9051).
1782        // It rebuilds via MoleculeBuilder::new() so groups are already empty; this
1783        // test guards against regressions where from_molecule() is accidentally used.
1784        use chematic_core::{AtomIdx, StereoGroup, StereoGroupKind};
1785        let mut mol = parse("[C@@H](F)(Cl)Br").unwrap();
1786        mol.add_stereo_group(StereoGroup::new(
1787            StereoGroupKind::Absolute,
1788            vec![AtomIdx(0)],
1789        ));
1790        assert_eq!(
1791            mol.stereo_groups().len(),
1792            1,
1793            "precondition: group was added"
1794        );
1795        let stripped = remove_stereo(&mol);
1796        assert_eq!(
1797            stripped.stereo_groups().len(),
1798            0,
1799            "remove_stereo must clear stereo groups"
1800        );
1801    }
1802
1803    #[test]
1804    fn clean_stereo_groups_drops_non_chiral_atoms() {
1805        // clean_stereo_groups filters out atoms with Chirality::None (RDKit PR #9051).
1806        use chematic_core::{AtomIdx, StereoGroup, StereoGroupKind};
1807        // atom 0 = CH3 (not chiral), atom 1 = @@ chiral center
1808        let mut mol = parse("C[C@@H](F)Cl").unwrap();
1809        mol.add_stereo_group(StereoGroup::new(
1810            StereoGroupKind::Absolute,
1811            vec![AtomIdx(0), AtomIdx(1)], // atom 0 is NOT chiral
1812        ));
1813        let cleaned = clean_stereo_groups(&mol);
1814        assert_eq!(
1815            cleaned.stereo_groups().len(),
1816            1,
1817            "group must survive with 1 atom"
1818        );
1819        assert_eq!(
1820            cleaned.stereo_groups()[0].atom_indices,
1821            vec![AtomIdx(1)],
1822            "only chiral atom must remain in group"
1823        );
1824    }
1825
1826    #[test]
1827    fn clean_stereo_groups_drops_empty_groups() {
1828        // A group with no chiral atoms at all must be removed entirely.
1829        use chematic_core::{AtomIdx, StereoGroup, StereoGroupKind};
1830        let mut mol = parse("CC").unwrap(); // no chiral atoms
1831        mol.add_stereo_group(StereoGroup::new(
1832            StereoGroupKind::Absolute,
1833            vec![AtomIdx(0)],
1834        ));
1835        let cleaned = clean_stereo_groups(&mol);
1836        assert_eq!(
1837            cleaned.stereo_groups().len(),
1838            0,
1839            "empty group must be removed"
1840        );
1841    }
1842
1843    #[test]
1844    fn clean_stereo_groups_preserves_valid_groups() {
1845        // A group whose atoms are all chiral must be kept intact.
1846        use chematic_core::{AtomIdx, StereoGroup, StereoGroupKind};
1847        let mut mol = parse("[C@@H](F)(Cl)Br").unwrap();
1848        mol.add_stereo_group(StereoGroup::new(
1849            StereoGroupKind::Absolute,
1850            vec![AtomIdx(0)],
1851        ));
1852        let cleaned = clean_stereo_groups(&mol);
1853        assert_eq!(
1854            cleaned.stereo_groups().len(),
1855            1,
1856            "valid group must be preserved"
1857        );
1858        assert_eq!(cleaned.stereo_groups()[0].atom_indices, vec![AtomIdx(0)]);
1859    }
1860
1861    #[test]
1862    fn normalize_groups_mixed_nitro_and_azide() {
1863        // Molecule with both nitro and azide groups
1864        let mol = parse("C[N+](=O)[O-].N[N+](=O)[O-]").unwrap();
1865        let result = normalize_groups(&mol);
1866
1867        // All atoms should be neutral
1868        let all_neutral = result.atoms().all(|(_, a)| a.charge == 0);
1869        assert!(all_neutral, "both nitro and azide should be neutralized");
1870    }
1871}