use std::collections::{HashMap, HashSet, VecDeque};
use chematic_core::{AtomIdx, BondIdx, Molecule};
use chematic_perception::find_sssr;
pub const BOND_LEN: f64 = 40.0;
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Point {
pub x: f64,
pub y: f64,
}
impl Point {
pub fn new(x: f64, y: f64) -> Self {
Self { x, y }
}
pub fn dist(&self, other: &Point) -> f64 {
let dx = self.x - other.x;
let dy = self.y - other.y;
(dx * dx + dy * dy).sqrt()
}
}
pub struct Layout {
pub coords: Vec<Point>,
}
impl Layout {
pub fn get(&self, idx: AtomIdx) -> Point {
self.coords[idx.0 as usize]
}
pub fn bounding_box(&self) -> (f64, f64, f64, f64) {
if self.coords.is_empty() {
return (0.0, 0.0, 0.0, 0.0);
}
self.coords.iter().fold(
(f64::MAX, f64::MAX, f64::MIN, f64::MIN),
|(min_x, min_y, max_x, max_y), p| {
(
min_x.min(p.x),
min_y.min(p.y),
max_x.max(p.x),
max_y.max(p.y),
)
},
)
}
}
pub fn compute_layout(mol: &Molecule) -> Layout {
let n = mol.atom_count();
if n == 0 {
return Layout { coords: Vec::new() };
}
if n == 1 {
return Layout {
coords: vec![Point::new(0.0, 0.0)],
};
}
let components = connected_components(mol);
let mut all_coords: Vec<Option<Point>> = vec![None; n];
let mut fragment_max_x = 0.0_f64;
for component_atoms in &components {
let component_set: HashSet<AtomIdx> = component_atoms.iter().copied().collect();
let mut placed: Vec<Option<Point>> = vec![None; n];
let ring_set = find_sssr(mol);
let rings: Vec<Vec<AtomIdx>> = ring_set
.rings()
.iter()
.filter(|ring| ring.iter().all(|a| component_set.contains(a)))
.cloned()
.collect();
let in_ring: HashSet<AtomIdx> = rings.iter().flat_map(|r| r.iter().copied()).collect();
let ring_systems = group_ring_systems(&rings);
for system in &ring_systems {
place_ring_system(mol, system, &mut placed);
}
place_chains(mol, &component_set, &in_ring, &mut placed);
let x_offset = if fragment_max_x == 0.0 {
0.0
} else {
fragment_max_x + 2.0 * BOND_LEN
};
let comp_min_x = component_atoms
.iter()
.filter_map(|&a| placed[a.0 as usize])
.map(|p| p.x)
.fold(f64::MAX, f64::min);
let shift = x_offset - comp_min_x;
for &a in component_atoms {
if let Some(p) = placed[a.0 as usize] {
let shifted = Point::new(p.x + shift, p.y);
all_coords[a.0 as usize] = Some(shifted);
if shifted.x > fragment_max_x {
fragment_max_x = shifted.x;
}
}
}
}
Layout {
coords: all_coords
.into_iter()
.map(|p| p.unwrap_or(Point::new(0.0, 0.0)))
.collect(),
}
}
fn connected_components(mol: &Molecule) -> Vec<Vec<AtomIdx>> {
let n = mol.atom_count();
let mut visited = vec![false; n];
let mut components: Vec<Vec<AtomIdx>> = Vec::new();
for start in 0..n {
if visited[start] {
continue;
}
let mut component = Vec::new();
let mut queue = VecDeque::new();
queue.push_back(AtomIdx(start as u32));
visited[start] = true;
while let Some(current) = queue.pop_front() {
component.push(current);
for (nb, _) in mol.neighbors(current) {
if !visited[nb.0 as usize] {
visited[nb.0 as usize] = true;
queue.push_back(nb);
}
}
}
components.push(component);
}
components
}
fn group_ring_systems(rings: &[Vec<AtomIdx>]) -> Vec<Vec<Vec<AtomIdx>>> {
if rings.is_empty() {
return Vec::new();
}
let n = rings.len();
let mut parent: Vec<usize> = (0..n).collect();
fn find(parent: &mut Vec<usize>, i: usize) -> usize {
if parent[i] != i {
parent[i] = find(parent, parent[i]);
}
parent[i]
}
fn union(parent: &mut Vec<usize>, a: usize, b: usize) {
let ra = find(parent, a);
let rb = find(parent, b);
if ra != rb {
parent[ra] = rb;
}
}
for (i, ring_i) in rings.iter().enumerate() {
let set_i: HashSet<AtomIdx> = ring_i.iter().copied().collect();
for (j, ring_j) in rings.iter().enumerate().skip(i + 1) {
if ring_j.iter().any(|a| set_i.contains(a)) {
union(&mut parent, i, j);
}
}
}
let mut groups: HashMap<usize, Vec<usize>> = HashMap::new();
for i in 0..n {
let root = find(&mut parent, i);
groups.entry(root).or_default().push(i);
}
groups
.values()
.map(|indices| indices.iter().map(|&i| rings[i].clone()).collect())
.collect()
}
fn place_ring_system(_mol: &Molecule, system: &[Vec<AtomIdx>], placed: &mut [Option<Point>]) {
if system.is_empty() {
return;
}
place_regular_ring(&system[0], placed);
let mut remaining: Vec<&Vec<AtomIdx>> = system[1..].iter().collect();
let mut iterations = 0;
while !remaining.is_empty() && iterations < remaining.len() * 2 {
iterations += 1;
let mut progressed = false;
remaining.retain(|ring| {
let already_placed: Vec<AtomIdx> = ring
.iter()
.copied()
.filter(|&a| placed[a.0 as usize].is_some())
.collect();
if already_placed.len() < 2 {
return true; }
let shared_edge = find_shared_edge(ring, placed);
let (anchor1, anchor2) = shared_edge.unwrap_or((already_placed[0], already_placed[1]));
let (Some(p1), Some(p2)) = (placed[anchor1.0 as usize], placed[anchor2.0 as usize])
else {
return true; };
place_ring_anchored(ring, anchor1, p1, anchor2, p2, placed);
progressed = true;
false });
if !progressed {
break;
}
}
for ring in &remaining {
place_regular_ring(ring, placed);
}
}
fn find_shared_edge(ring: &[AtomIdx], placed: &[Option<Point>]) -> Option<(AtomIdx, AtomIdx)> {
let n = ring.len();
ring.windows(2)
.map(|w| (w[0], w[1]))
.chain(std::iter::once((ring[n - 1], ring[0])))
.find(|&(a, b)| placed[a.0 as usize].is_some() && placed[b.0 as usize].is_some())
}
fn place_regular_ring(ring: &[AtomIdx], placed: &mut [Option<Point>]) {
let n = ring.len();
if n == 0 {
return;
}
let radius = ring_radius(n);
let start_angle = std::f64::consts::FRAC_PI_2;
for (i, &atom) in ring.iter().enumerate() {
if placed[atom.0 as usize].is_none() {
let angle = start_angle - (2.0 * std::f64::consts::PI * i as f64) / n as f64;
let x = radius * angle.cos();
let y = -radius * angle.sin(); placed[atom.0 as usize] = Some(Point::new(x, y));
}
}
}
fn place_ring_anchored(
ring: &[AtomIdx],
anchor1: AtomIdx,
p1: Point,
anchor2: AtomIdx,
p2: Point,
placed: &mut [Option<Point>],
) {
let n = ring.len();
let radius = ring_radius(n);
let idx1 = ring.iter().position(|&a| a == anchor1).unwrap_or(0);
let idx2 = ring.iter().position(|&a| a == anchor2).unwrap_or(1);
let mid = Point::new((p1.x + p2.x) / 2.0, (p1.y + p2.y) / 2.0);
let dx = p2.x - p1.x;
let dy = p2.y - p1.y;
let edge_len = (dx * dx + dy * dy).sqrt();
if edge_len < 1e-10 {
return;
}
let perp_x = -dy / edge_len;
let perp_y = dx / edge_len;
let apothem = radius * (std::f64::consts::PI / n as f64).cos();
let existing_center = {
let pts: Vec<Point> = ring.iter().filter_map(|&a| placed[a.0 as usize]).collect();
if pts.is_empty() {
mid
} else {
let cx = pts.iter().map(|p| p.x).sum::<f64>() / pts.len() as f64;
let cy = pts.iter().map(|p| p.y).sum::<f64>() / pts.len() as f64;
Point::new(cx, cy)
}
};
let cand1 = Point::new(mid.x + perp_x * apothem, mid.y + perp_y * apothem);
let cand2 = Point::new(mid.x - perp_x * apothem, mid.y - perp_y * apothem);
let new_center = if cand1.dist(&existing_center) > cand2.dist(&existing_center) {
cand1
} else {
cand2
};
let angle_to_a1 = (p1.y - new_center.y).atan2(p1.x - new_center.x);
let steps_forward = (idx2 + n - idx1) % n; let angle_to_a2 = (p2.y - new_center.y).atan2(p2.x - new_center.x);
let mut delta = angle_to_a2 - angle_to_a1;
while delta > std::f64::consts::PI {
delta -= 2.0 * std::f64::consts::PI;
}
while delta < -std::f64::consts::PI {
delta += 2.0 * std::f64::consts::PI;
}
let angle_step = if steps_forward > 0 {
if delta >= 0.0 {
2.0 * std::f64::consts::PI / n as f64
} else {
-(2.0 * std::f64::consts::PI / n as f64)
}
} else {
2.0 * std::f64::consts::PI / n as f64
};
for step in 0..n {
let ring_idx = (idx1 + step) % n;
let atom = ring[ring_idx];
if placed[atom.0 as usize].is_some() {
continue;
}
let angle = angle_to_a1 + step as f64 * angle_step;
let x = new_center.x + radius * angle.cos();
let y = new_center.y + radius * angle.sin();
placed[atom.0 as usize] = Some(Point::new(x, y));
}
}
fn ring_radius(n: usize) -> f64 {
if n < 3 {
return BOND_LEN;
}
BOND_LEN / (2.0 * (std::f64::consts::PI / n as f64).sin())
}
fn place_chains(
mol: &Molecule,
component: &HashSet<AtomIdx>,
in_ring: &HashSet<AtomIdx>,
placed: &mut [Option<Point>],
) {
let ring_atoms: Vec<AtomIdx> = component
.iter()
.copied()
.filter(|a| in_ring.contains(a) && placed[a.0 as usize].is_some())
.collect();
for start in &ring_atoms {
let unplaced_neighbors: Vec<AtomIdx> = mol
.neighbors(*start)
.map(|(nb, _)| nb)
.filter(|nb| component.contains(nb) && placed[nb.0 as usize].is_none())
.collect();
for nb in unplaced_neighbors {
if placed[nb.0 as usize].is_some() {
continue;
}
let dir = best_outgoing_direction(*start, placed, mol, component);
dfs_zigzag(mol, nb, *start, dir, placed, component);
}
}
let unplaced: Vec<AtomIdx> = component
.iter()
.copied()
.filter(|a| placed[a.0 as usize].is_none())
.collect();
if unplaced.is_empty() {
return;
}
let start = unplaced
.iter()
.copied()
.find(|&a| mol.degree(a) <= 1)
.unwrap_or(unplaced[0]);
placed[start.0 as usize] = Some(Point::new(0.0, 0.0));
let init_dir = 0.0_f64;
let neighbors: Vec<AtomIdx> = mol
.neighbors(start)
.map(|(nb, _)| nb)
.filter(|nb| component.contains(nb) && placed[nb.0 as usize].is_none())
.collect();
for (i, nb) in neighbors.into_iter().enumerate() {
if placed[nb.0 as usize].is_some() {
continue;
}
let dir = if i == 0 {
init_dir
} else {
init_dir + std::f64::consts::PI
};
dfs_zigzag(mol, nb, start, dir, placed, component);
}
let still_unplaced: Vec<AtomIdx> = component
.iter()
.copied()
.filter(|a| placed[a.0 as usize].is_none())
.collect();
let mut x = 0.0;
for atom in still_unplaced {
x += BOND_LEN;
placed[atom.0 as usize] = Some(Point::new(x, 0.0));
}
}
fn best_outgoing_direction(
atom: AtomIdx,
placed: &[Option<Point>],
mol: &Molecule,
component: &HashSet<AtomIdx>,
) -> f64 {
let Some(origin) = placed[atom.0 as usize] else {
return 0.0;
};
let used_angles: Vec<f64> = mol
.neighbors(atom)
.filter(|(nb, _)| component.contains(nb) && placed[nb.0 as usize].is_some())
.map(|(nb, _)| {
let p = placed[nb.0 as usize].unwrap();
(p.y - origin.y).atan2(p.x - origin.x)
})
.collect();
if used_angles.is_empty() {
return 0.0;
}
let candidates: Vec<f64> = (0..6)
.map(|i| i as f64 * std::f64::consts::PI / 3.0)
.collect();
candidates
.into_iter()
.max_by(|&a, &b| {
let min_sep_a = min_angle_separation(a, &used_angles);
let min_sep_b = min_angle_separation(b, &used_angles);
min_sep_a.partial_cmp(&min_sep_b).unwrap()
})
.unwrap_or(0.0)
}
fn min_angle_separation(angle: f64, used: &[f64]) -> f64 {
used.iter()
.map(|&u| {
let diff = (angle - u).abs();
if diff > std::f64::consts::PI {
2.0 * std::f64::consts::PI - diff
} else {
diff
}
})
.fold(f64::MAX, f64::min)
}
fn dfs_zigzag(
mol: &Molecule,
start_atom: AtomIdx,
start_parent: AtomIdx,
start_dir: f64,
placed: &mut [Option<Point>],
component: &HashSet<AtomIdx>,
) {
let deflections = [-std::f64::consts::PI / 6.0, std::f64::consts::PI / 6.0];
let mut stack: Vec<(AtomIdx, AtomIdx, f64)> = vec![(start_atom, start_parent, start_dir)];
while let Some((atom, parent, dir)) = stack.pop() {
if placed[atom.0 as usize].is_some() {
continue;
}
let parent_pos = match placed[parent.0 as usize] {
Some(p) => p,
None => continue,
};
let x = parent_pos.x + BOND_LEN * dir.cos();
let y = parent_pos.y + BOND_LEN * dir.sin();
placed[atom.0 as usize] = Some(Point::new(x, y));
let unplaced: Vec<AtomIdx> = mol
.neighbors(atom)
.map(|(nb, _)| nb)
.filter(|&nb| {
nb != parent && component.contains(&nb) && placed[nb.0 as usize].is_none()
})
.collect();
for (i, nb) in unplaced.into_iter().enumerate().rev() {
stack.push((nb, atom, dir + deflections[i % 2]));
}
}
}
pub fn suggest_bond_direction(mol: &Molecule, atom: AtomIdx, layout: &Layout) -> f64 {
use std::f64::consts::PI;
let origin = layout.get(atom);
let used_angles: Vec<f64> = mol
.neighbors(atom)
.filter(|(nb, _)| (nb.0 as usize) < layout.coords.len())
.map(|(nb, _)| {
let p = layout.get(nb);
(p.y - origin.y).atan2(p.x - origin.x)
})
.collect();
if used_angles.is_empty() {
return 0.0;
}
let mut candidates: Vec<f64> = (0..12).map(|i| i as f64 * PI / 6.0).collect();
for &a in &used_angles {
candidates.push(a + PI);
candidates.push(a + PI - PI / 6.0);
candidates.push(a + PI + PI / 6.0);
candidates.push(a + 2.0 * PI / 3.0);
candidates.push(a - 2.0 * PI / 3.0);
}
candidates
.into_iter()
.max_by(|&ca, &cb| {
let sa = min_angle_separation(ca, &used_angles);
let sb = min_angle_separation(cb, &used_angles);
sa.partial_cmp(&sb).unwrap()
})
.unwrap_or(0.0)
}
pub fn detect_crossings(layout: &Layout, mol: &Molecule) -> Vec<(BondIdx, BondIdx)> {
let bonds: Vec<(BondIdx, (Point, Point))> = mol
.bonds()
.map(|(bidx, bond)| {
let p1 = layout.get(bond.atom1);
let p2 = layout.get(bond.atom2);
(bidx, (p1, p2))
})
.collect();
let mut crossings = Vec::new();
for i in 0..bonds.len() {
for j in (i + 1)..bonds.len() {
let (bidx_i, (p1_i, p2_i)) = bonds[i];
let (bidx_j, (p1_j, p2_j)) = bonds[j];
let a1_i = mol.bond(bidx_i).atom1;
let a2_i = mol.bond(bidx_i).atom2;
let a1_j = mol.bond(bidx_j).atom1;
let a2_j = mol.bond(bidx_j).atom2;
if a1_i == a1_j || a1_i == a2_j || a2_i == a1_j || a2_i == a2_j {
continue;
}
if segments_intersect(p1_i, p2_i, p1_j, p2_j) {
crossings.push((bidx_i, bidx_j));
}
}
}
crossings
}
fn segments_intersect(a: Point, b: Point, c: Point, d: Point) -> bool {
let ccw = |p1: Point, p2: Point, p3: Point| -> bool {
(p3.y - p1.y) * (p2.x - p1.x) > (p2.y - p1.y) * (p3.x - p1.x)
};
ccw(a, c, d) != ccw(b, c, d) && ccw(a, b, c) != ccw(a, b, d)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ring_radius_hexagon() {
let r = ring_radius(6);
let expected = BOND_LEN; assert!((r - expected).abs() < 1e-9, "hexagon radius = {}", r);
}
fn make_layout(coords: &[(f64, f64)]) -> Layout {
Layout {
coords: coords.iter().map(|&(x, y)| Point { x, y }).collect(),
}
}
#[test]
fn test_suggest_direction_no_neighbors() {
use chematic_smiles::parse;
let mol = parse("C").unwrap();
let layout = make_layout(&[(0.0, 0.0)]);
let dir = suggest_bond_direction(&mol, AtomIdx(0), &layout);
assert!((dir).abs() < 1e-9 || (dir - 2.0 * std::f64::consts::PI).abs() < 1e-9);
}
#[test]
fn test_suggest_direction_single_bond_avoids_existing() {
use chematic_smiles::parse;
let mol = parse("CC").unwrap();
let layout = make_layout(&[(0.0, 0.0), (BOND_LEN, 0.0)]);
let dir = suggest_bond_direction(&mol, AtomIdx(0), &layout);
let sep = min_angle_separation(dir, &[0.0_f64]);
assert!(
sep >= std::f64::consts::PI / 2.0,
"suggested direction {dir:.3} should be ≥90° from existing bond, sep={sep:.3}"
);
}
#[test]
fn test_suggest_direction_two_bonds_finds_gap() {
use chematic_smiles::parse;
let mol = parse("CCC").unwrap();
let layout = make_layout(&[(-BOND_LEN, 0.0), (0.0, 0.0), (BOND_LEN, 0.0)]);
let dir = suggest_bond_direction(&mol, AtomIdx(1), &layout);
let sep = {
let used = [0.0_f64, std::f64::consts::PI];
min_angle_separation(dir, &used)
};
assert!(
sep >= std::f64::consts::PI / 2.0 - 1e-6,
"center atom: suggested direction {dir:.3} should be ~90° from both bonds, sep={sep:.3}"
);
}
#[test]
fn test_suggest_direction_prefers_sp2_for_aromatic_ring() {
use chematic_smiles::parse;
let mol = parse("c1ccccc1").unwrap();
let layout = compute_layout(&mol);
let dir = suggest_bond_direction(&mol, AtomIdx(0), &layout);
let p0 = layout.get(AtomIdx(0));
let used: Vec<f64> = mol
.neighbors(AtomIdx(0))
.map(|(nb, _)| {
let p = layout.get(nb);
(p.y - p0.y).atan2(p.x - p0.x)
})
.collect();
let sep = min_angle_separation(dir, &used);
assert!(
sep >= std::f64::consts::PI / 3.0 - 1e-6,
"benzene exit direction should be ≥60° from ring bonds, sep={sep:.3}"
);
}
}