1use serde::{Deserialize, Serialize};
7
8use crate::scene::MoleculeIn;
9
10#[cfg(feature = "codegen")]
11use schemars::JsonSchema;
12#[cfg(feature = "codegen")]
13use ts_rs::TS;
14
15pub const MIN_MCS_ATOMS: u32 = 3;
17
18#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
20#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
21#[cfg_attr(feature = "codegen", ts(export))]
22pub struct AlignOpts {
23 #[serde(default, skip_serializing_if = "Option::is_none")]
25 #[cfg_attr(feature = "codegen", ts(optional))]
26 pub atom_map: Option<Vec<(u32, u32)>>,
27 #[serde(default, skip_serializing_if = "Option::is_none")]
29 #[cfg_attr(feature = "codegen", ts(optional))]
30 pub min_atoms: Option<u32>,
31}
32
33#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
35#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
36#[cfg_attr(feature = "codegen", ts(export))]
37pub struct MolTemplate {
38 pub id: String,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
41 #[cfg_attr(feature = "codegen", ts(optional))]
42 pub smiles: Option<String>,
43 #[serde(default, skip_serializing_if = "Option::is_none")]
44 #[cfg_attr(feature = "codegen", ts(optional))]
45 pub cxsmiles: Option<String>,
46 #[serde(default, skip_serializing_if = "Option::is_none")]
47 #[cfg_attr(feature = "codegen", ts(optional))]
48 pub molfile: Option<String>,
49 #[serde(default, skip_serializing_if = "Option::is_none")]
51 #[cfg_attr(feature = "codegen", ts(optional))]
52 pub align: Option<AlignOpts>,
53 #[serde(default, skip_serializing_if = "Vec::is_empty")]
55 pub template_for: Vec<MolTemplate>,
56}
57
58impl MolTemplate {
59 pub fn validate_structure(&self) -> Result<(), String> {
61 let n = [&self.smiles, &self.cxsmiles, &self.molfile]
62 .iter()
63 .filter(|s| s.as_ref().is_some_and(|t| !t.trim().is_empty()))
64 .count();
65 if n == 1 {
66 Ok(())
67 } else {
68 Err(format!(
69 "MolTemplate {}: need exactly one of smiles/cxsmiles/molfile",
70 self.id
71 ))
72 }
73 }
74
75 pub fn source(&self) -> Result<&str, String> {
76 self.validate_structure()?;
77 if let Some(s) = self.smiles.as_deref().filter(|t| !t.trim().is_empty()) {
78 return Ok(s);
79 }
80 if let Some(s) = self.cxsmiles.as_deref().filter(|t| !t.trim().is_empty()) {
81 return Ok(s);
82 }
83 if let Some(s) = self.molfile.as_deref().filter(|t| !t.trim().is_empty()) {
84 return Ok(s);
85 }
86 Err(format!("MolTemplate {}: empty structure", self.id))
87 }
88}
89
90#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
92#[serde(tag = "type", rename_all = "snake_case")]
93#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
94#[cfg_attr(feature = "codegen", ts(export))]
95pub enum EdgeTask {
96 CoordGen {
97 #[serde(default)]
98 roots: Vec<MolTemplate>,
99 },
100}
101
102#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
104#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
105#[cfg_attr(feature = "codegen", ts(export))]
106pub struct EdgePlan {
107 pub version: u32,
108 #[serde(default)]
109 pub tasks: Vec<EdgeTask>,
110}
111
112impl EdgePlan {
113 pub fn new_v1(tasks: Vec<EdgeTask>) -> Self {
114 Self {
115 version: 1,
116 tasks,
117 }
118 }
119
120 pub fn validate(&self) -> Result<(), String> {
121 if self.version != 1 {
122 return Err(format!("unsupported EdgePlan version {}", self.version));
123 }
124 let mut seen = std::collections::HashSet::new();
125 for task in &self.tasks {
126 match task {
127 EdgeTask::CoordGen { roots } => {
128 for root in roots {
129 validate_tree(root, true, &mut seen)?;
130 }
131 }
132 }
133 }
134 Ok(())
135 }
136}
137
138fn validate_tree(
139 node: &MolTemplate,
140 is_root: bool,
141 seen: &mut std::collections::HashSet<String>,
142) -> Result<(), String> {
143 if !seen.insert(node.id.clone()) {
144 return Err(format!("duplicate MolTemplate id {}", node.id));
145 }
146 node.validate_structure()?;
147 if is_root && node.align.is_some() {
148 return Err(format!(
149 "MolTemplate {}: roots must have align=null",
150 node.id
151 ));
152 }
153 if !is_root && node.align.is_none() {
154 }
156 for child in &node.template_for {
157 validate_tree(child, false, seen)?;
158 }
159 Ok(())
160}
161
162#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
164#[serde(rename_all = "snake_case")]
165#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
166#[cfg_attr(feature = "codegen", ts(export))]
167pub enum CoordMethod {
168 Free,
169 AtomMap,
170 Mcs,
171 None,
173}
174
175#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
177#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
178#[cfg_attr(feature = "codegen", ts(export))]
179pub struct CoordGenMoleculeResult {
180 pub id: String,
181 pub ok: bool,
183 pub method: CoordMethod,
184 #[serde(default, skip_serializing_if = "Option::is_none")]
185 #[cfg_attr(feature = "codegen", ts(optional))]
186 pub used_map: Option<Vec<(u32, u32)>>,
187 #[serde(default, skip_serializing_if = "Option::is_none")]
189 #[cfg_attr(feature = "codegen", ts(optional))]
190 pub molecule: Option<MoleculeIn>,
191 #[serde(default, skip_serializing_if = "Option::is_none")]
192 #[cfg_attr(feature = "codegen", ts(optional))]
193 pub error: Option<String>,
194}
195
196#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
197#[serde(tag = "type", rename_all = "snake_case")]
198#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
199#[cfg_attr(feature = "codegen", ts(export))]
200pub enum EdgeTaskResult {
201 CoordGen {
202 ok: bool,
203 molecules: Vec<CoordGenMoleculeResult>,
204 },
205}
206
207#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
208#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
209#[cfg_attr(feature = "codegen", ts(export))]
210pub struct EdgeResult {
211 pub version: u32,
212 #[serde(default)]
213 pub results: Vec<EdgeTaskResult>,
214}
215
216impl EdgeResult {
217 pub fn new_v1(results: Vec<EdgeTaskResult>) -> Self {
218 Self {
219 version: 1,
220 results,
221 }
222 }
223
224 pub fn molecules(&self) -> Vec<&CoordGenMoleculeResult> {
226 let mut out = Vec::new();
227 for r in &self.results {
228 match r {
229 EdgeTaskResult::CoordGen { molecules, .. } => out.extend(molecules.iter()),
230 }
231 }
232 out
233 }
234}
235
236#[cfg(test)]
237mod tests {
238 use super::*;
239
240 fn sample_plan() -> EdgePlan {
241 EdgePlan::new_v1(vec![EdgeTask::CoordGen {
242 roots: vec![MolTemplate {
243 id: "m_0".into(),
244 smiles: Some("c1ccccc1".into()),
245 cxsmiles: None,
246 molfile: None,
247 align: None,
248 template_for: vec![MolTemplate {
249 id: "m_1".into(),
250 smiles: Some("Cc1ccccc1".into()),
251 cxsmiles: None,
252 molfile: None,
253 align: Some(AlignOpts {
254 atom_map: Some(vec![
255 (1, 0),
256 (2, 1),
257 (3, 2),
258 (4, 3),
259 (5, 4),
260 (6, 5),
261 ]),
262 min_atoms: None,
263 }),
264 template_for: vec![],
265 }],
266 }],
267 }])
268 }
269
270 #[test]
271 fn plan_json_roundtrip() {
272 let plan = sample_plan();
273 plan.validate().unwrap();
274 let s = serde_json::to_string(&plan).unwrap();
275 let back: EdgePlan = serde_json::from_str(&s).unwrap();
276 assert_eq!(back, plan);
277 }
278
279 #[test]
280 fn reject_duplicate_ids() {
281 let mut plan = sample_plan();
282 match &mut plan.tasks[0] {
283 EdgeTask::CoordGen { roots } => {
284 roots[0].template_for[0].id = "m_0".into();
285 }
286 }
287 assert!(plan.validate().unwrap_err().contains("duplicate"));
288 }
289
290 #[test]
291 fn reject_root_with_align() {
292 let plan = EdgePlan::new_v1(vec![EdgeTask::CoordGen {
293 roots: vec![MolTemplate {
294 id: "m_0".into(),
295 smiles: Some("CCO".into()),
296 cxsmiles: None,
297 molfile: None,
298 align: Some(AlignOpts::default()),
299 template_for: vec![],
300 }],
301 }]);
302 assert!(plan.validate().unwrap_err().contains("align=null"));
303 }
304
305 #[test]
306 fn fallback_result_lists_all_ids() {
307 let result = EdgeResult::new_v1(vec![EdgeTaskResult::CoordGen {
308 ok: true,
309 molecules: vec![
310 CoordGenMoleculeResult {
311 id: "m_0".into(),
312 ok: true,
313 method: CoordMethod::Free,
314 used_map: None,
315 molecule: None,
316 error: None,
317 },
318 CoordGenMoleculeResult {
319 id: "m_1".into(),
320 ok: true,
321 method: CoordMethod::None,
322 used_map: None,
323 molecule: None,
324 error: Some("align failed; fell back".into()),
325 },
326 ],
327 }]);
328 let ids: Vec<_> = result.molecules().iter().map(|m| m.id.as_str()).collect();
329 assert_eq!(ids, ["m_0", "m_1"]);
330 assert!(matches!(
331 &result.results[0],
332 EdgeTaskResult::CoordGen { ok: true, .. }
333 ));
334 }
335}