1use std::path::PathBuf;
25
26use serde::{Deserialize, Serialize};
27
28const EDF_HEADER_BLOCK: usize = 256;
30
31#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
38#[serde(rename_all = "lowercase")]
39pub enum SampleDtype {
40 I16,
42 I32,
44 I64,
46}
47
48impl SampleDtype {
49 pub fn width(self) -> usize {
51 match self {
52 SampleDtype::I16 => 2,
53 SampleDtype::I32 => 4,
54 SampleDtype::I64 => 8,
55 }
56 }
57
58 pub fn as_token(self) -> &'static str {
61 match self {
62 SampleDtype::I16 => "i16",
63 SampleDtype::I32 => "i32",
64 SampleDtype::I64 => "i64",
65 }
66 }
67}
68
69impl Default for SampleDtype {
70 fn default() -> Self {
73 SampleDtype::I32
74 }
75}
76
77pub struct ScratchDir {
84 pub path: PathBuf,
86}
87
88impl ScratchDir {
89 pub fn new(tag: &str) -> std::io::Result<Self> {
93 let nanos = std::time::SystemTime::now()
94 .duration_since(std::time::UNIX_EPOCH)
95 .map(|d| d.as_nanos())
96 .unwrap_or(0);
97 use std::sync::atomic::{AtomicU64, Ordering};
98 static SEQ: AtomicU64 = AtomicU64::new(0);
99 let seq = SEQ.fetch_add(1, Ordering::Relaxed);
100 let name = format!("lqs_{}_{}_{}_{}", tag, std::process::id(), nanos, seq);
101 let path = std::env::temp_dir().join(name);
102 std::fs::create_dir_all(&path)?;
103 Ok(Self { path })
104 }
105
106 pub fn join(&self, name: &str) -> PathBuf {
108 self.path.join(name)
109 }
110}
111
112impl Drop for ScratchDir {
113 fn drop(&mut self) {
114 let _ = std::fs::remove_dir_all(&self.path);
116 }
117}
118
119fn edf_field(out: &mut Vec<u8>, value: &str, width: usize) -> Result<(), ()> {
122 let bytes = value.as_bytes();
123 if bytes.len() > width {
124 return Err(());
125 }
126 out.extend_from_slice(bytes);
127 out.resize(out.len() + (width - bytes.len()), b' ');
128 Ok(())
129}
130
131fn format_edf_number(x: f64) -> String {
136 if x.fract() == 0.0 && x.abs() < 1e15 {
137 format!("{}", x as i64)
138 } else {
139 let s = format!("{x:.6}");
140 let trimmed = s.trim_end_matches('0').trim_end_matches('.');
141 trimmed.to_string()
142 }
143}
144
145pub fn write_edf_bytes(signal: &[Vec<i64>], fs: f64) -> Option<Vec<u8>> {
154 let ns = signal.len();
155 if ns == 0 {
156 return None;
157 }
158 let spr = signal[0].len();
160 if spr == 0 || signal.iter().any(|c| c.len() != spr) {
161 return None;
162 }
163 if signal
165 .iter()
166 .flat_map(|c| c.iter())
167 .any(|&s| s < i16::MIN as i64 || s > i16::MAX as i64)
168 {
169 return None;
170 }
171 if !fs.is_finite() || fs <= 0.0 {
172 return None;
173 }
174
175 let record_duration = spr as f64 / fs;
178 let dur_str = format_edf_number(record_duration);
179 if dur_str.len() > 8 {
180 return None;
181 }
182
183 let header_bytes = EDF_HEADER_BLOCK + ns * EDF_HEADER_BLOCK;
184 let header_str = header_bytes.to_string();
185 if header_str.len() > 8 || ns.to_string().len() > 4 || spr.to_string().len() > 8 {
186 return None;
187 }
188
189 let mut buf = Vec::with_capacity(header_bytes + ns * spr * 2);
190
191 edf_field(&mut buf, "0", 8).ok()?; edf_field(&mut buf, "OpenECS X X X", 80).ok()?; edf_field(&mut buf, "Startdate X", 80).ok()?; edf_field(&mut buf, "01.01.26", 8).ok()?; edf_field(&mut buf, "00.00.00", 8).ok()?; edf_field(&mut buf, &header_str, 8).ok()?; edf_field(&mut buf, "", 44).ok()?; edf_field(&mut buf, "1", 8).ok()?; edf_field(&mut buf, &dur_str, 8).ok()?; edf_field(&mut buf, &ns.to_string(), 4).ok()?; for i in 0..ns {
205 edf_field(&mut buf, &format!("ch{i}"), 16).ok()?; }
207 for _ in 0..ns {
208 edf_field(&mut buf, "AgAgCl", 80).ok()?; }
210 for _ in 0..ns {
211 edf_field(&mut buf, "uV", 8).ok()?; }
213 for _ in 0..ns {
214 edf_field(&mut buf, "-32768", 8).ok()?; }
216 for _ in 0..ns {
217 edf_field(&mut buf, "32767", 8).ok()?; }
219 for _ in 0..ns {
220 edf_field(&mut buf, "-32768", 8).ok()?; }
222 for _ in 0..ns {
223 edf_field(&mut buf, "32767", 8).ok()?; }
225 for _ in 0..ns {
226 edf_field(&mut buf, "", 80).ok()?; }
228 for _ in 0..ns {
229 edf_field(&mut buf, &spr.to_string(), 8).ok()?; }
231 for _ in 0..ns {
232 edf_field(&mut buf, "", 32).ok()?; }
234 debug_assert_eq!(buf.len(), header_bytes, "EDF header block size mismatch");
235
236 for chan in signal {
238 for &s in chan {
239 buf.extend_from_slice(&(s as i16).to_le_bytes());
240 }
241 }
242
243 Some(buf)
244}
245
246pub fn reshape_channel_major(
256 raw: &[u8],
257 n_chan: usize,
258 per_chan: usize,
259 dtype: SampleDtype,
260) -> Option<Vec<Vec<i64>>> {
261 let width = dtype.width();
262 let total = n_chan.checked_mul(per_chan)?.checked_mul(width)?;
263 if raw.len() != total {
264 return None;
265 }
266
267 let mut out = Vec::with_capacity(n_chan);
268 let mut pos = 0usize;
269 for _ in 0..n_chan {
270 let mut chan = Vec::with_capacity(per_chan);
271 for _ in 0..per_chan {
272 let v = match dtype {
273 SampleDtype::I16 => {
274 i16::from_le_bytes([raw[pos], raw[pos + 1]]) as i64
275 }
276 SampleDtype::I32 => {
277 i32::from_le_bytes([raw[pos], raw[pos + 1], raw[pos + 2], raw[pos + 3]])
278 as i64
279 }
280 SampleDtype::I64 => i64::from_le_bytes([
281 raw[pos],
282 raw[pos + 1],
283 raw[pos + 2],
284 raw[pos + 3],
285 raw[pos + 4],
286 raw[pos + 5],
287 raw[pos + 6],
288 raw[pos + 7],
289 ]),
290 };
291 chan.push(v);
292 pos += width;
293 }
294 out.push(chan);
295 }
296 Some(out)
297}
298
299#[cfg(test)]
300mod tests {
301 use super::*;
302
303 fn fixture() -> Vec<Vec<i64>> {
304 vec![vec![0, 1, -1, 1000, -1000], vec![5, 6, 7, 8, 9]]
305 }
306
307 #[test]
308 fn scratchdir_is_unique_and_cleans_up() {
309 let a = ScratchDir::new("t").expect("scratch a");
310 let b = ScratchDir::new("t").expect("scratch b");
311 assert_ne!(a.path, b.path, "two scratch dirs must not collide");
312 let p = a.path.clone();
313 assert!(p.is_dir());
314 drop(a);
315 assert!(!p.exists(), "scratch dir removed on drop");
316 drop(b);
317 }
318
319 #[test]
320 fn write_edf_rejects_unsupported_shapes() {
321 assert!(write_edf_bytes(&[], 256.0).is_none());
322 assert!(write_edf_bytes(&[vec![]], 256.0).is_none());
323 assert!(write_edf_bytes(&[vec![1, 2], vec![1]], 256.0).is_none());
324 assert!(write_edf_bytes(&[vec![i64::MAX]], 256.0).is_none());
325 assert!(write_edf_bytes(&[vec![1, 2]], 0.0).is_none());
326 let edf = write_edf_bytes(&fixture(), 256.0).expect("valid fixture -> EDF");
327 let ns = 2usize;
328 let spr = 5usize;
329 assert_eq!(edf.len(), EDF_HEADER_BLOCK * (1 + ns) + ns * spr * 2);
330 assert_eq!(&edf[..8], b"0 ", "EDF version field");
331 }
332
333 #[test]
334 fn reshape_round_trips_each_dtype() {
335 for dt in [SampleDtype::I16, SampleDtype::I32, SampleDtype::I64] {
336 let sig = vec![vec![0i64, 1, -1, 100], vec![-100, 2, 3, 4]];
337 let mut raw = Vec::new();
339 for chan in &sig {
340 for &s in chan {
341 match dt {
342 SampleDtype::I16 => raw.extend_from_slice(&(s as i16).to_le_bytes()),
343 SampleDtype::I32 => raw.extend_from_slice(&(s as i32).to_le_bytes()),
344 SampleDtype::I64 => raw.extend_from_slice(&s.to_le_bytes()),
345 }
346 }
347 }
348 let back = reshape_channel_major(&raw, 2, 4, dt).expect("valid stream");
349 assert_eq!(back, sig, "round trip for {dt:?}");
350 }
351 }
352
353 #[test]
354 fn reshape_rejects_wrong_length() {
355 let raw = vec![0u8; 30];
357 assert!(reshape_channel_major(&raw, 2, 4, SampleDtype::I32).is_none());
358 assert!(reshape_channel_major(&[], usize::MAX, usize::MAX, SampleDtype::I64).is_none());
360 }
361
362 #[test]
363 fn dtype_width_and_token() {
364 assert_eq!(SampleDtype::I16.width(), 2);
365 assert_eq!(SampleDtype::I32.width(), 4);
366 assert_eq!(SampleDtype::I64.width(), 8);
367 assert_eq!(SampleDtype::default(), SampleDtype::I32);
368 assert_eq!(SampleDtype::I16.as_token(), "i16");
369 }
370}