dspk-pipeline 0.0.10

//! Everything decided before the device is touched: the stages, the opened input, the channels and
//! range, the mode and halos, the window schedule. [`plan`] reports it ([`Plan`]); the executor
//! (`run.rs`) runs it.

use std::path::{Path, PathBuf};

use dspk_core::compute::device_elements;
use dspk_core::{ChunkSchedule, DspError, DspResult, RecordingSource, SampleFormat, SignalUnit};
use dspk_io::MergedRecording;
use dspk_ops::core::EdgeMode;
use dspk_ops::filter::FilterMode;
use dspk_ops::resampler::poly::RESAMPLE_POLY_HALF_LEN_PER_RATE;
use dspk_ops::resampler::{resample_poly_len, FirWindow, ResampleFilter};
use dspk_ops::{Pipeline, PipelineStage};

use crate::spec::{ExecMode, Input, Node, OutputFormat, PipelineSpec, Upload};
use crate::PipelineError;

/// How the chain runs.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExecutedMode {
    /// Windows in order with filter state carried: equal to processing the whole range at once.
    Stream,
    /// Independent windows with `left` / `right` halos of context.
    Windows { left: usize, right: usize },
}

/// What running one input would do, without touching the device.
#[derive(Debug, Clone)]
pub struct Plan {
    pub input: String,
    pub output: PathBuf,
    pub mode: ExecutedMode,
    /// `true` when stored integers would be uploaded and scaled on the device.
    pub stored_upload: bool,
    /// Channels read (and processed).
    pub channels: usize,
    /// Channels written (fewer when a remap before write drops some).
    pub output_channels: usize,
    /// Input samples per channel processed (`stop - start`).
    pub input_samples: u64,
    pub input_rate_hz: f64,
    /// Output samples per channel and sample rate (differ from the input's after `resample`).
    pub output_samples: u64,
    pub output_rate_hz: f64,
    pub windows: usize,
    /// Samples per channel of the largest window, halos included.
    pub max_window_samples: usize,
    /// Bytes read from the input and uploaded, over all windows (halos included).
    pub bytes_read: u64,
    pub bytes_uploaded: u64,
}

/// A trailing `resample` node, reduced to coprime factors.
pub(crate) struct Resample {
    pub(crate) up: usize,
    pub(crate) down: usize,
    pub(crate) filter: ResampleFilter,
    pub(crate) edge: EdgeMode,
    /// Input samples the resampling filter reaches on each side of an output sample.
    pub(crate) reach: usize,
}

impl Resample {
    fn new(up: usize, down: usize, beta: f64, edge: EdgeMode) -> Self {
        let g = gcd(up, down);
        let (up, down) = (up / g, down / g);
        let half_len = RESAMPLE_POLY_HALF_LEN_PER_RATE * up.max(down);
        Self { up, down, filter: ResampleFilter::Window(FirWindow::Kaiser { beta }), edge, reach: half_len.div_ceil(up) + 1 }
    }
}

/// One input, ready to run.
pub(crate) struct Prepared {
    pub(crate) source: Box<dyn RecordingSource>,
    pub(crate) pipeline: Pipeline,
    /// The pipeline's stages with the node each comes from (branches of one spec share the
    /// nodes, and so the stages, of their common start).
    pub(crate) stages: Vec<(String, dspk_ops::pipeline::PipelineStage)>,
    pub(crate) resample: Option<Resample>,
    pub(crate) channels: Vec<usize>,
    /// The output channels' names (the input's, or a remap's).
    pub(crate) names: Vec<String>,
    /// The pipeline output's channels written, in order (a remap after the last processing node);
    /// `None`: all.
    pub(crate) keep: Option<Vec<usize>>,
    pub(crate) range: std::ops::Range<u64>,
    pub(crate) unit: SignalUnit,
    pub(crate) sample_rate: f64,
    pub(crate) stored_format: SampleFormat,
    pub(crate) stored_upload: bool,
    pub(crate) mode: ExecutedMode,
    pub(crate) schedule: ChunkSchedule,
    pub(crate) out_rate: f64,
    pub(crate) out_samples: u64,
    pub(crate) output: PathBuf,
    pub(crate) format: OutputFormat,
    /// Chunking and compression of Zarr output; `parallel` also applies to raw output.
    pub(crate) zarr: dspk_io::ZarrOptions,
    pub(crate) input_label: String,
    /// Whitening stages still to fit, by their index in `stages` (identities until then).
    pub(crate) fits: Vec<(usize, crate::WhitenFit)>,
    /// A `crossings` node ending the chain: its events are the output (`format = "events"`).
    pub(crate) crossings: Option<dspk_ops::crossings::CrossingOptions>,
}

impl Prepared {
    pub(crate) fn plan(&self) -> Plan {
        let nch = self.channels.len() as u64;
        let read: u64 = self.schedule.windows().iter().map(|w| w.read_len() as u64).sum::<u64>() * nch;
        let per_sample_up = if self.stored_upload { self.stored_format.bytes() } else { 4 } as u64;
        Plan {
            input: self.input_label.clone(),
            output: self.output.clone(),
            mode: self.mode,
            stored_upload: self.stored_upload,
            channels: self.channels.len(),
            output_channels: self.names.len(),
            input_samples: self.range.end - self.range.start,
            input_rate_hz: self.sample_rate,
            output_samples: self.out_samples,
            output_rate_hz: self.out_rate,
            windows: self.schedule.len(),
            max_window_samples: self.schedule.max_read_samples(),
            bytes_read: read * self.stored_format.bytes() as u64,
            bytes_uploaded: read * per_sample_up,
        }
    }
}

/// What [`crate::run`] would do for each input of `spec`, opening the inputs (metadata reads only)
/// with `dspk_io::open` and never touching the device.
pub fn plan(spec: &PipelineSpec, base_dir: &Path) -> Result<Vec<Plan>, PipelineError> {
    plan_with(spec, base_dir, default_open)
}

/// [`plan`] with recordings opened by `open`.
pub fn plan_with(
    spec: &PipelineSpec,
    base_dir: &Path,
    open: impl Fn(&Path, Option<&str>) -> DspResult<Box<dyn RecordingSource>>,
) -> Result<Vec<Plan>, PipelineError> {
    spec.expand_inputs(base_dir)?.iter().map(|one| prepare(one, base_dir, &open).map(|p| p.plan())).collect()
}

/// Opens a recording with `dspk_io` (generic formats): `source` picks a source within the file.
pub fn default_open(path: &Path, source: Option<&str>) -> DspResult<Box<dyn RecordingSource>> {
    match source {
        Some(id) => dspk_io::open_source(path, id),
        None => dspk_io::open(path),
    }
}

/// One input of an expanded spec (a `path`, or a single merge), checked and scheduled.
pub(crate) fn prepare(
    spec: &PipelineSpec,
    base_dir: &Path,
    open: &impl Fn(&Path, Option<&str>) -> DspResult<Box<dyn RecordingSource>>,
) -> Result<Prepared, PipelineError> {
    let chain = spec.chain()?;
    let Node::Read { path: input, inputs, source, channels: picked, start, stop } = chain.read.1 else {
        unreachable!("chain starts with read")
    };
    let Node::Write { path: output, format, compression, chunk_samples: zarr_chunk, parallel } = chain.write.1 else { unreachable!("chain ends with write") };
    let parallel = parallel.unwrap_or(dspk_io::DEFAULT_PARALLEL_WRITES);
    let zarr = match format {
        OutputFormat::Zarr => {
            let mut options = dspk_io::ZarrOptions { parallel, ..Default::default() };
            if let Some(c) = compression {
                options.compression = (*c).into();
            }
            if let Some(n) = zarr_chunk {
                if *n == 0 {
                    return Err(PipelineError::Invalid("write: chunk_samples must be at least 1".into()));
                }
                options.chunk_samples = *n;
            }
            options
        }
        OutputFormat::Raw | OutputFormat::Events if compression.is_some() || zarr_chunk.is_some() => {
            return Err(PipelineError::Invalid("write: compression and chunk_samples apply to zarr output only".into()));
        }
        OutputFormat::Raw | OutputFormat::Events => dspk_io::ZarrOptions { parallel, ..Default::default() },
    };

    // Remaps by position: right after read (fold into the read: may drop), between processing
    // steps (a reorder stage on the device: keeps every channel), after the last one (a gather
    // before the download: may drop)
    let last_processing = chain.stages.iter().rposition(|(_, n)| !matches!(n, Node::Remap { .. }));
    let mut lead = Vec::new();
    let mut tail = Vec::new();
    // The pipeline's steps in order: a stage, or a middle remap (its order is set once the
    // channel count is known)
    enum Step {
        Stage(PipelineStage),
        Reorder(dspk_core::ChannelMap),
        Whiten(crate::WhitenFit),
    }
    let mut steps: Vec<(String, Step)> = Vec::with_capacity(chain.stages.len());
    let mut resample = None;
    // (level, release, direction, min duration ms, refractory ms) of a final `crossings` node
    let mut crossings_node = None;
    for (i, (id, node)) in chain.stages.iter().enumerate() {
        let stage = node.stage().map_err(|e| PipelineError::Invalid(format!("node {id:?}: {e}")))?;
        match (node, stage) {
            (Node::Remap { .. }, _) => {
                let map = node.channel_map(Some(base_dir)).map_err(|e| PipelineError::Invalid(format!("node {id:?}: {e}")))?.expect("a remap has a map");
                match last_processing {
                    Some(last) if i > last => tail.push(map),
                    Some(_) if steps.is_empty() && resample.is_none() => lead.push(map),
                    None => lead.push(map),
                    Some(_) => steps.push((id.to_string(), Step::Reorder(map))),
                }
            }
            (Node::Resample { up, down, beta, edge }, _) => {
                if last_processing != Some(i) {
                    return Err(PipelineError::Invalid(format!("node {id:?}: resample must be the last processing node before write (only remaps may follow)")));
                }
                resample = Some(Resample::new(*up, *down, *beta, (*edge).into()));
            }
            (Node::Crossings { level, release, direction, min_duration_ms, refractory_ms }, _) => {
                if last_processing != Some(i) {
                    return Err(PipelineError::Invalid(format!("node {id:?}: crossings must be the last processing node before write (only remaps may follow)")));
                }
                crossings_node = Some((*level, *release, *direction, *min_duration_ms, *refractory_ms));
            }
            (Node::Whiten { .. }, None) => steps.push((id.to_string(), Step::Whiten(node.whiten_fit().expect("a whiten node")))),
            (_, Some(stage)) => steps.push((id.to_string(), Step::Stage(stage))),
            (_, None) => unreachable!("every other processing node has a stage"),
        }
    }

    let (source, input_label): (Box<dyn RecordingSource>, String) = match (input, inputs.as_deref()) {
        (Some(path), None) => (open(&base_dir.join(path), source.as_deref())?, path.display().to_string()),
        (None, Some([Input::Merge { merge }])) => {
            let parts = merge
                .iter()
                .map(|p| Ok((open(&base_dir.join(&p.path), p.source.as_deref())?, p.channels.clone())))
                .collect::<DspResult<Vec<_>>>()?;
            let label = merge.iter().map(|p| p.path.display().to_string()).collect::<Vec<_>>().join(" + ");
            (Box::new(MergedRecording::new(parts)?), label)
        }
        _ => unreachable!("expand_inputs leaves one path or one merge"),
    };
    let info = source.info();
    let (total, sample_rate, stored_format) = (info.samples, info.sample_rate_hz(), info.format);
    let mut channels = select_channels(picked.as_deref(), info.channel_count())?;
    let remap_err = |e: DspError| PipelineError::Invalid(format!("remap: {e}"));
    let mut names: Vec<String> = channels.iter().map(|&c| info.channels[c].name.clone()).collect();
    // `names` through `map` (0-based over the current channels): the names it gives, else the
    // current ones
    let rename = |names: &[String], map: &dspk_core::ChannelMap, order: &[usize]| -> Vec<String> {
        order.iter().zip(map.names()).map(|(&c, n)| n.unwrap_or_else(|| names[c].clone())).collect()
    };
    if !lead.is_empty() {
        let mut map = dspk_core::ChannelMap::from_order(&channels);
        for r in &lead {
            map = map.then(r).map_err(remap_err)?;
        }
        channels = map.resolve(info.channel_count()).map_err(remap_err)?;
        let all: Vec<String> = info.channels.iter().map(|c| c.name.clone()).collect();
        names = rename(&all, &map, &channels);
    }
    let nch = channels.len();
    let mut stages = Vec::with_capacity(steps.len());
    let mut fits = Vec::new();
    for (id, step) in steps {
        let stage = match step {
            Step::Stage(stage) => stage,
            Step::Whiten(fit) => {
                fits.push((stages.len(), fit));
                PipelineStage::SpatialWhitening(crate::unfitted_whitening(nch))
            }
            Step::Reorder(map) => {
                let order = map.resolve(nch).map_err(|e| PipelineError::Invalid(format!("node {id:?}: {e}")))?;
                if order.len() != nch {
                    return Err(PipelineError::Invalid(format!(
                        "node {id:?}: a remap between processing nodes keeps every channel ({nch}; it reorders them): drop channels in a remap right after read or right before write"
                    )));
                }
                names = rename(&names, &map, &order);
                PipelineStage::ReorderChannels { order }
            }
        };
        stages.push((id, stage));
    }
    let keep = if tail.is_empty() {
        None
    } else {
        let mut map = dspk_core::ChannelMap::from_order(&(0..nch).collect::<Vec<_>>());
        for r in &tail {
            map = map.then(r).map_err(remap_err)?;
        }
        let order = map.resolve(nch).map_err(remap_err)?;
        names = rename(&names, &map, &order);
        Some(order)
    };
    match (&crossings_node, format) {
        (Some(_), OutputFormat::Events) | (None, OutputFormat::Raw | OutputFormat::Zarr) => {}
        (Some(_), _) => return Err(PipelineError::Invalid("a crossings node gives events: write them with format = \"events\"".into())),
        (None, _) => return Err(PipelineError::Invalid("write: format = \"events\" needs a crossings node before it".into())),
    }
    if crossings_node.is_some() && resample.is_some() {
        return Err(PipelineError::Invalid("crossings after resample: detect on the signal at its own rate".into()));
    }
    let stage_ids = stages.clone();
    let pipeline = Pipeline::with_stages(stages.into_iter().map(|(_, s)| s).collect());
    let range = select_range(*start, *stop, total)?;
    let unit = info.channels[channels[0]].unit.clone();
    let to_samples = |ms: f64| (ms * 1e-3 * sample_rate).round() as u64;
    let crossings = crossings_node.map(|(level, release, direction, min_ms, refr_ms)| dspk_ops::crossings::CrossingOptions {
        levels: vec![level],
        releases: release.map(|r| vec![r]).unwrap_or_default(),
        direction: direction.into(),
        min_duration_samples: to_samples(min_ms),
        refractory_samples: to_samples(refr_ms),
    });
    pipeline.validate(sample_rate).map_err(|e| PipelineError::Invalid(e.to_string()))?;

    let mut mode = resolve_mode(spec.execution.mode, &pipeline, sample_rate)?;
    let stored_upload = resolve_upload(spec.execution.upload, stored_format)?;
    let mut chunk = spec.execution.chunk_samples.max(1);
    let (mut out_rate, mut out_samples) = (sample_rate, range.end - range.start);
    if let Some(r) = &resample {
        // Independent windows aligned to `down`, with the stages' context plus the resampling
        // filter's reach: each output sample sees the input a whole-recording run would
        if spec.execution.mode == ExecMode::Stream {
            return Err(PipelineError::Invalid("execution.mode = \"stream\" is not possible with resample: use \"windows\" or \"auto\"".into()));
        }
        if range.start % r.down as u64 != 0 {
            return Err(PipelineError::Invalid(format!("read: start {} must be a multiple of the resample's down ({})", range.start, r.down)));
        }
        let (left, right) = pipeline.settling(sample_rate).map_err(|e| PipelineError::Invalid(e.to_string()))?;
        let align = |x: usize| x.div_ceil(r.down) * r.down;
        mode = ExecutedMode::Windows { left: align(left + r.reach), right: align(right + r.reach) };
        chunk = align(chunk);
        out_rate = sample_rate * r.up as f64 / r.down as f64;
        out_samples = resample_poly_len((range.end - range.start) as usize, r.up, r.down) as u64;
    }
    let (left, right) = match mode {
        ExecutedMode::Stream => (0, 0),
        ExecutedMode::Windows { left, right } => (left, right),
    };
    let schedule = ChunkSchedule::new(range.clone(), chunk as u64, left as u64, right as u64, total);
    device_elements("pipeline window [channels, samples]", &[channels.len(), schedule.max_read_samples()])?;

    Ok(Prepared {
        source,
        pipeline,
        stages: stage_ids,
        resample,
        channels,
        names,
        keep,
        range,
        unit,
        sample_rate,
        stored_format,
        stored_upload,
        mode,
        schedule,
        out_rate,
        out_samples,
        output: base_dir.join(output),
        format: *format,
        zarr,
        input_label,
        fits,
        crossings,
    })
}

fn gcd(a: usize, b: usize) -> usize {
    if b == 0 { a } else { gcd(b, a % b) }
}

/// `picked` (in its order, no repeats), or every channel.
fn select_channels(picked: Option<&[usize]>, count: usize) -> Result<Vec<usize>, PipelineError> {
    let Some(picked) = picked else { return Ok((0..count).collect()) };
    if picked.is_empty() {
        return Err(PipelineError::Invalid("read: channels is empty".into()));
    }
    let mut seen = vec![false; count];
    for &c in picked {
        if c >= count {
            return Err(PipelineError::Invalid(format!("read: channel {c} does not exist (the recording has {count})")));
        }
        if std::mem::replace(&mut seen[c], true) {
            return Err(PipelineError::Invalid(format!("read: channel {c} is listed twice")));
        }
    }
    Ok(picked.to_vec())
}

fn select_range(start: Option<u64>, stop: Option<u64>, total: u64) -> Result<std::ops::Range<u64>, PipelineError> {
    let (start, stop) = (start.unwrap_or(0), stop.unwrap_or(total));
    if start >= stop || stop > total {
        return Err(PipelineError::Invalid(format!("read: samples {start}..{stop} are not a non-empty range within 0..{total}")));
    }
    Ok(start..stop)
}

/// Stored upload works for integer storage; float32 gains nothing; float64 has no device unpacking.
fn resolve_upload(requested: Upload, format: SampleFormat) -> Result<bool, PipelineError> {
    let integer = !matches!(format, SampleFormat::F32 | SampleFormat::F64);
    match requested {
        Upload::Auto => Ok(integer),
        Upload::Scaled => Ok(false),
        Upload::Stored if format == SampleFormat::F64 => {
            Err(PipelineError::Invalid("execution.upload = \"stored\": float64 recordings cannot be scaled on the device".into()))
        }
        Upload::Stored => Ok(true),
    }
}

/// `stream` needs every filter forward (zero-phase needs future samples) and no stencil stage
/// (median, Gaussian, Teager–Kaiser, FIR read neighbours across window edges and carry no state;
/// for FIR, window mode with its halo is exact).
fn resolve_mode(requested: ExecMode, pipeline: &Pipeline, sample_rate: f64) -> Result<ExecutedMode, PipelineError> {
    let blocker = pipeline.stages().iter().find_map(|stage| match stage {
        PipelineStage::Filter(spec) if spec.mode == FilterMode::ForwardBackward => Some("a zero-phase (forward_backward) filter"),
        PipelineStage::Median { .. } => Some("a median stage"),
        PipelineStage::Fir { .. } => Some("a fir stage"),
        PipelineStage::GaussianSmooth { .. } => Some("a gaussian_smooth stage"),
        PipelineStage::TeagerKaiser { .. } => Some("a teager_kaiser stage"),
        _ => None,
    });
    let windows = || {
        let (left, right) = pipeline.settling(sample_rate).map_err(|e| PipelineError::Invalid(e.to_string()))?;
        Ok(ExecutedMode::Windows { left, right })
    };
    match (requested, blocker) {
        (ExecMode::Stream, Some(what)) => Err(PipelineError::Invalid(format!(
            "execution.mode = \"stream\" is not possible with {what}: use \"windows\" or \"auto\""
        ))),
        (ExecMode::Stream, None) | (ExecMode::Auto, None) => Ok(ExecutedMode::Stream),
        (ExecMode::Windows, _) | (ExecMode::Auto, Some(_)) => windows(),
    }
}

impl Prepared {
    /// Fits the whitening stages on the device (each on the signal of the stages before it, the
    /// ones fitted before included), so the pipeline is ready to run.
    pub(crate) fn fit(&mut self, target: dspk_core::ComputeTarget) -> Result<(), PipelineError> {
        if self.fits.is_empty() {
            return Ok(());
        }
        for (at, fit) in std::mem::take(&mut self.fits) {
            let before: Vec<PipelineStage> = self.stages[..at].iter().map(|(_, s)| s.clone()).collect();
            let fitted = crate::fit_whitening(target, self.source.as_ref(), &self.channels, &before, &fit)?;
            self.stages[at].1 = PipelineStage::SpatialWhitening(fitted);
        }
        self.pipeline = Pipeline::with_stages(self.stages.iter().map(|(_, s)| s.clone()).collect());
        Ok(())
    }
}