1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::{DataError, Result};
6use crate::ids::{SampleId, SourceId};
7use crate::model::PresenceMask;
8
9#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum AlignmentMode {
12 #[default]
13 Inner,
14 Left,
15 Outer,
16}
17
18#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
19pub struct AlignmentPolicy {
20 pub mode: AlignmentMode,
21}
22
23impl Default for AlignmentPolicy {
24 fn default() -> Self {
25 Self {
26 mode: AlignmentMode::Inner,
27 }
28 }
29}
30
31#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
32pub struct SourceSampleSet {
33 pub source_id: SourceId,
34 pub sample_ids: Vec<SampleId>,
35}
36
37#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
38pub struct SampleAlignmentPlan {
39 pub mode: AlignmentMode,
40 pub sample_ids: Vec<SampleId>,
41 pub masks: Vec<PresenceMask>,
42}
43
44impl SampleAlignmentPlan {
45 pub fn validate(&self) -> Result<()> {
46 if self.sample_ids.is_empty() {
47 return Err(DataError::Validation(
48 "alignment plan contains no samples".to_string(),
49 ));
50 }
51 let mut samples = BTreeSet::new();
52 for sample_id in &self.sample_ids {
53 if !samples.insert(sample_id) {
54 return Err(DataError::Validation(format!(
55 "alignment plan contains duplicate sample `{sample_id}`"
56 )));
57 }
58 }
59 if self.masks.is_empty() {
60 return Err(DataError::Validation(
61 "alignment plan contains no presence masks".to_string(),
62 ));
63 }
64 let mut sources = BTreeSet::new();
65 for mask in &self.masks {
66 mask.validate()?;
67 if mask.sample_ids != self.sample_ids {
68 return Err(DataError::Validation(format!(
69 "presence mask for `{}` does not use alignment sample order",
70 mask.source_id
71 )));
72 }
73 if !sources.insert(&mask.source_id) {
74 return Err(DataError::Validation(format!(
75 "alignment plan contains duplicate source mask `{}`",
76 mask.source_id
77 )));
78 }
79 }
80 Ok(())
81 }
82}
83
84pub fn build_sample_alignment_plan(
85 sources: &[SourceSampleSet],
86 policy: &AlignmentPolicy,
87) -> Result<SampleAlignmentPlan> {
88 if sources.is_empty() {
89 return Err(DataError::Validation(
90 "alignment requires at least one source".to_string(),
91 ));
92 }
93 let mut source_ids = BTreeSet::new();
94 let mut samples_by_source = Vec::with_capacity(sources.len());
95 for source in sources {
96 if !source_ids.insert(&source.source_id) {
97 return Err(DataError::Validation(format!(
98 "alignment contains duplicate source `{}`",
99 source.source_id
100 )));
101 }
102 if source.sample_ids.is_empty() {
103 return Err(DataError::Validation(format!(
104 "alignment source `{}` contains no samples",
105 source.source_id
106 )));
107 }
108 let mut samples = BTreeSet::new();
109 for sample_id in &source.sample_ids {
110 if !samples.insert(sample_id) {
111 return Err(DataError::Validation(format!(
112 "alignment source `{}` contains duplicate sample `{sample_id}`",
113 source.source_id
114 )));
115 }
116 }
117 samples_by_source.push(samples);
118 }
119
120 let sample_ids = match policy.mode {
121 AlignmentMode::Inner => sources[0]
122 .sample_ids
123 .iter()
124 .filter(|sample_id| {
125 samples_by_source
126 .iter()
127 .all(|samples| samples.contains(*sample_id))
128 })
129 .cloned()
130 .collect::<Vec<_>>(),
131 AlignmentMode::Left => sources[0].sample_ids.clone(),
132 AlignmentMode::Outer => {
133 let mut seen = BTreeSet::new();
134 let mut ordered = Vec::new();
135 for source in sources {
136 for sample_id in &source.sample_ids {
137 if seen.insert(sample_id) {
138 ordered.push(sample_id.clone());
139 }
140 }
141 }
142 ordered
143 }
144 };
145
146 if sample_ids.is_empty() {
147 return Err(DataError::Validation(format!(
148 "{:?} alignment produced no shared samples",
149 policy.mode
150 )));
151 }
152
153 let masks = sources
154 .iter()
155 .zip(samples_by_source.iter())
156 .map(|(source, samples)| PresenceMask {
157 sample_ids: sample_ids.clone(),
158 source_id: source.source_id.clone(),
159 present: sample_ids
160 .iter()
161 .map(|sample_id| samples.contains(sample_id))
162 .collect(),
163 })
164 .collect::<Vec<_>>();
165 let plan = SampleAlignmentPlan {
166 mode: policy.mode,
167 sample_ids,
168 masks,
169 };
170 plan.validate()?;
171 Ok(plan)
172}
173
174pub fn alignment_mode_from_fusion(value: Option<&serde_json::Value>) -> Result<AlignmentMode> {
175 let Some(value) = value else {
176 return Ok(AlignmentMode::Inner);
177 };
178 let Some(mode) = value.get("alignment") else {
179 return Ok(AlignmentMode::Inner);
180 };
181 serde_json::from_value(mode.clone())
182 .map_err(|error| DataError::Validation(format!("invalid fusion alignment policy: {error}")))
183}
184
185pub fn alignment_metadata(
186 inputs: Vec<String>,
187 output: String,
188 mode: AlignmentMode,
189) -> Result<BTreeMap<String, serde_json::Value>> {
190 Ok(BTreeMap::from([
191 (
192 "inputs".to_string(),
193 serde_json::Value::Array(inputs.into_iter().map(serde_json::Value::String).collect()),
194 ),
195 ("output".to_string(), serde_json::Value::String(output)),
196 ("alignment".to_string(), serde_json::to_value(mode)?),
197 ]))
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203
204 fn source(source_id: &str, sample_ids: &[&str]) -> SourceSampleSet {
205 SourceSampleSet {
206 source_id: SourceId::new(source_id).unwrap(),
207 sample_ids: sample_ids
208 .iter()
209 .map(|sample_id| SampleId::new(*sample_id).unwrap())
210 .collect(),
211 }
212 }
213
214 fn samples(plan: &SampleAlignmentPlan) -> Vec<String> {
215 plan.sample_ids.iter().map(ToString::to_string).collect()
216 }
217
218 #[test]
219 fn inner_alignment_keeps_first_source_order_for_shared_samples() {
220 let plan = build_sample_alignment_plan(
221 &[
222 source("nir", &["S002", "S001", "S003"]),
223 source("chem", &["S001", "S003"]),
224 ],
225 &AlignmentPolicy {
226 mode: AlignmentMode::Inner,
227 },
228 )
229 .unwrap();
230
231 assert_eq!(samples(&plan), vec!["S001", "S003"]);
232 assert_eq!(plan.masks[0].present, vec![true, true]);
233 assert_eq!(plan.masks[1].present, vec![true, true]);
234 }
235
236 #[test]
237 fn left_alignment_preserves_left_samples_and_marks_missing_sources() {
238 let plan = build_sample_alignment_plan(
239 &[
240 source("nir", &["S002", "S001", "S003"]),
241 source("chem", &["S001", "S003"]),
242 ],
243 &AlignmentPolicy {
244 mode: AlignmentMode::Left,
245 },
246 )
247 .unwrap();
248
249 assert_eq!(samples(&plan), vec!["S002", "S001", "S003"]);
250 assert_eq!(plan.masks[0].present, vec![true, true, true]);
251 assert_eq!(plan.masks[1].present, vec![false, true, true]);
252 }
253
254 #[test]
255 fn outer_alignment_appends_new_samples_in_source_order() {
256 let plan = build_sample_alignment_plan(
257 &[
258 source("nir", &["S002", "S001"]),
259 source("chem", &["S003", "S001"]),
260 source("image", &["S004", "S002"]),
261 ],
262 &AlignmentPolicy {
263 mode: AlignmentMode::Outer,
264 },
265 )
266 .unwrap();
267
268 assert_eq!(samples(&plan), vec!["S002", "S001", "S003", "S004"]);
269 assert_eq!(plan.masks[0].present, vec![true, true, false, false]);
270 assert_eq!(plan.masks[1].present, vec![false, true, true, false]);
271 assert_eq!(plan.masks[2].present, vec![true, false, false, true]);
272 }
273
274 #[test]
275 fn alignment_refuses_duplicate_samples_per_source() {
276 let err = build_sample_alignment_plan(
277 &[source("nir", &["S001", "S001"])],
278 &AlignmentPolicy::default(),
279 )
280 .unwrap_err();
281
282 assert!(err.to_string().contains("duplicate sample"));
283 }
284}