deimos 0.16.2

//! A lookup-table sequence machine that follows a set procedure during
//! each sequence, and transitions between sequences based on set criteria.

use core::f64;
use std::fmt::Write;
use std::{
    collections::{BTreeSet, HashSet},
    path::Path,
};

use serde::{Deserialize, Serialize};
use serde_json;

pub type StateName = String;

#[cfg(feature = "python")]
use pyo3::prelude::*;

use super::*;

mod lookup;
mod sequence;
mod transition;

pub use lookup::{InterpMethod, SequenceLookup};
pub use sequence::Sequence;
pub use transition::{ThreshOp, Timeout, Transition};

#[derive(Default, Debug)]
struct ExecutionState {
    /// Time in current sequence's sequence.
    /// Starts at the first time in the sequence's lookup table.
    pub sequence_time_s: f64,

    /// Name of the current operating sequence
    pub current_sequence: String,

    // Values provided by calc orchestrator during init
    /// Lookup map for channel names to indices, required for evaluating
    /// transition criteria
    pub input_index_map: BTreeMap<String, usize>,

    /// Timestep
    pub dt_s: f64,

    /// Indices of input calc/channel names
    pub input_indices: Vec<usize>,

    /// Where to write the calc outputs in the calc tape
    pub output_range: Range<usize>,
}

/// Sequence entrypoint and transition criteria for the SequenceMachine.
#[derive(Default, Debug, Serialize, Deserialize)]
pub struct MachineCfg {
    // User inputs
    /// Whether to dispatch outputs
    pub save_outputs: bool,

    /// Name of Sequence which is the entrypoint for the machine
    pub entry: String,

    /// Whether to reload from a folder at this relative path from the op dir during init
    pub link_folder: Option<String>,

    /// Timeout behavior for each sequence
    pub timeouts: BTreeMap<String, Timeout>,

    /// Early transition criteria for each sequence
    pub transitions: BTreeMap<String, BTreeMap<String, Vec<Transition>>>,
}

/// A lookup-table sequence machine that follows a set procedure during
/// each sequence, and transitions between sequences based on set criteria.
///
/// Unlike most calcs, the names of the inputs and outputs of this calc
/// are not known at compile-time, and are assembled from inputs instead.
#[derive(Serialize, Deserialize, Debug)]
#[cfg_attr(feature = "python", pyclass)]
pub struct SequenceMachine {
    /// State transition criteria and other configuration
    cfg: MachineCfg,

    /// All the lookup sequence sequences of the machine, including their
    /// transition criteria.
    ///
    /// All sequences must have the same outputs so that no values
    /// are ever left dangling, and all values must be defined
    /// at the first timestep of each sequence.
    ///
    /// The inputs to the machine are the sum of all the inputs
    /// required by each sequence.
    sequences: BTreeMap<String, Sequence>,

    /// Current execution state of the SequenceMachine
    /// including sequence time and per-run configuration.
    #[serde(skip)]
    execution_state: ExecutionState,
}

impl Default for SequenceMachine {
    fn default() -> Self {
        Self {
            cfg: MachineCfg {
                entry: "Placeholder".into(),
                ..Default::default()
            },
            sequences: BTreeMap::from([("Placeholder".into(), Sequence::default())]),
            execution_state: ExecutionState::default(),
        }
    }
}

impl SequenceMachine {
    /// Store and validate a new SequenceMachine
    pub fn new(
        cfg: MachineCfg,
        sequences: BTreeMap<String, Sequence>,
    ) -> Result<Box<Self>, String> {
        // These will be set during init.
        // Use default indices that will cause an error on the first call if not initialized properly
        let input_indices = Vec::new();
        let output_range = usize::MAX..usize::MAX;
        let entry = cfg.entry.to_owned();

        let machine = Self {
            cfg,
            sequences,
            execution_state: ExecutionState {
                sequence_time_s: f64::NAN,
                current_sequence: entry,
                input_index_map: BTreeMap::new(),
                dt_s: f64::NAN,
                input_indices,
                output_range,
            },
        };

        machine.validate()?;

        Ok(Box::new(machine))
    }

    /// Check the validity of sequences, transitions, timeouts, etc
    fn validate(&self) -> Result<(), String> {
        // Validate individual sequences and lookups
        for seq in self.sequences.values() {
            seq.validate()?;
        }

        // Validate machine-level configuration
        let seq_names: HashSet<String> = self.sequences.keys().cloned().collect();

        // Make sure all the timeouts are present and there aren't any extras
        let timeout_seq_names: HashSet<String> = self.cfg.timeouts.keys().cloned().collect();
        if seq_names != timeout_seq_names {
            return Err(format!(
                "Timeouts do not match sequences. Sequence names that are present in sequences but not timeouts: `{:?}`. Sequence names that are present in timeouts but not sequences: {:?}",
                seq_names.difference(&timeout_seq_names),
                timeout_seq_names.difference(&seq_names)
            ));
        }

        // Make sure all the transitions are present and there aren't any extras
        let transition_seq_names: HashSet<String> = self.cfg.timeouts.keys().cloned().collect();
        if seq_names != transition_seq_names {
            return Err(format!(
                "Transitions do not match sequences. Sequence names that are present in sequences but not timeouts: `{:?}`. Sequence names that are present in transitions but not sequences: {:?}",
                seq_names.difference(&transition_seq_names),
                transition_seq_names.difference(&seq_names)
            ));
        }

        // Make sure all the transitions' target sequences refer to real targets
        for (seq, transitions) in self.cfg.transitions.iter() {
            for target_sequence in transitions.keys() {
                if !seq_names.contains(target_sequence) {
                    return Err(format!(
                        "Sequence `{seq}` has transition target sequence `{target_sequence}` which does not exist."
                    ));
                }
            }
        }

        Ok(())
    }

    #[cfg(feature = "python")]
    fn add_transition(
        &mut self,
        source_sequence: String,
        target_sequence: String,
        transition: Transition,
    ) -> Result<(), String> {
        if !self.sequences.contains_key(&source_sequence) {
            return Err(format!("Unknown source sequence: {source_sequence}"));
        }
        if !self.sequences.contains_key(&target_sequence) {
            return Err(format!("Unknown target sequence: {target_sequence}"));
        }

        self.cfg
            .transitions
            .entry(source_sequence)
            .or_default()
            .entry(target_sequence)
            .or_default()
            .push(transition);

        Ok(())
    }

    /// Get a reference to the sequence indicated in execution_state.current_sequence
    fn current_sequence(&self) -> &Sequence {
        &self.sequences[&self.execution_state.current_sequence]
    }

    /// Get a reference to the entrypoint sequence
    fn entry_sequence(&self) -> &Sequence {
        &self.sequences[&self.cfg.entry]
    }

    /// Set next target sequence and reset sequence time to the initial time for that sequence.
    fn transition(&mut self, target_sequence: String) {
        self.execution_state.current_sequence = target_sequence;
        self.execution_state.sequence_time_s = self.current_sequence().get_start_time_s();
    }

    /// Check each sequence transition criterion and set the next sequence if needed.
    /// If multiple transition criteria are met at the same time, the first
    /// one in the list will be prioritized.
    fn check_transitions(&mut self, sequence_time_s: f64, tape: &[f64]) -> Result<(), String> {
        let sequence_name = &self.execution_state.current_sequence;

        // Check for timeout
        if sequence_time_s > self.current_sequence().get_end_time_s() {
            return match &self.cfg.timeouts[sequence_name] {
                Timeout::Transition(target_sequence) => {
                    // info!("Transition `{target_sequence}` due to timeout");
                    self.transition(target_sequence.clone());
                    Ok(())
                }
                Timeout::Loop => {
                    // info!("Looping sequence `{sequence_name}` due to timeout");
                    self.transition(sequence_name.clone());
                    Ok(())
                }
            };
        }

        // Check other criteria
        for (target_sequence, criteria) in self.cfg.transitions[sequence_name].iter() {
            // Check whether this each criterion has been met
            for criterion in criteria {
                let should_transition = match criterion {
                    Transition::ConstantThresh(channel, op, thresh) => {
                        let i = self.execution_state.input_index_map[channel];
                        let v = tape[i];

                        op.eval(v, *thresh)
                    }
                    Transition::ChannelThresh(val_channel, op, thresh_channel) => {
                        let ival = self.execution_state.input_index_map[val_channel];
                        let ithresh = self.execution_state.input_index_map[thresh_channel];
                        let v = tape[ival];
                        let thresh = tape[ithresh];

                        op.eval(v, thresh)
                    }
                    Transition::LookupThresh(channel, op, lookup) => {
                        let i = self.execution_state.input_index_map[channel];
                        let v = tape[i];
                        let thresh = lookup.eval(sequence_time_s);

                        op.eval(v, thresh)
                    }
                };

                // If a sequence transition has been triggered, update the execution sequence
                // to the start of the next sequence.
                if should_transition {
                    // info!("Transition `{target_sequence}` due to {criterion:?}");
                    self.transition(target_sequence.clone());
                    return Ok(());
                }
            }
        }

        // No transition criteria were met; stay the course
        Ok(())
    }

    /// Read a configuration json and sequence CSV files from a folder.
    /// The folder must contain one json representing a [MachineCfg] and
    /// some number of CSV files each representing a [Sequence].
    pub fn load_folder(path: &dyn AsRef<Path>) -> Result<Box<Self>, String> {
        let dir = std::fs::read_dir(path)
            .map_err(|e| format!("Unable to read items in folder {:?}: {e}", path.as_ref()))?;

        let mut csv_files = Vec::new();
        let mut json_files = Vec::new();
        for e in dir.flatten() {
            let path = e.path();
            if path.is_file() {
                match path.extension() {
                    Some(ext) if ext.to_ascii_lowercase().to_str() == Some("csv") => {
                        csv_files.push(path)
                    }
                    Some(ext) if ext.to_ascii_lowercase().to_str() == Some("json") => {
                        json_files.push(path)
                    }
                    _ => {}
                }
            }
        }

        // Make sure there is exactly one json file
        if json_files.is_empty() {
            return Err("Did not find configuration json file".to_string());
        }

        if json_files.len() > 1 {
            return Err(format!("Found multiple config json files: {json_files:?}"));
        }

        // Load config
        let json_file = &json_files[0];
        let json_str = std::fs::read_to_string(json_file)
            .map_err(|e| format!("Failed to read config json: {e}"))?;
        let cfg: MachineCfg = serde_json::from_str(&json_str)
            .map_err(|e| format!("Failed to parse config json: {e}"))?;

        // Load sequences
        let mut sequences = BTreeMap::new();
        for fp in csv_files {
            // unwrap: ok because we already checked that this is a file with an extension
            let name = fp
                .file_stem()
                .ok_or_else(|| "Filename missing".to_string())?
                .to_str()
                .ok_or_else(|| "Filename is not valid unicode".to_string())?
                .to_owned();
            let seq: Sequence = Sequence::from_csv_file(&fp)?;
            sequences.insert(name, seq);
        }

        Self::new(cfg, sequences)
    }

    /// Save a configuration json and sequence CSV files to a folder.
    pub fn save_folder(&self, path: &dyn AsRef<Path>) -> Result<(), String> {
        let dir = path.as_ref();
        std::fs::create_dir_all(dir)
            .map_err(|e| format!("Unable to create folder {:?}: {e}", dir))?;

        let cfg_path = dir.join("cfg.json");
        let cfg_json = serde_json::to_string_pretty(&self.cfg)
            .map_err(|e| format!("Failed to serialize config json: {e}"))?;
        std::fs::write(&cfg_path, cfg_json)
            .map_err(|e| format!("Failed to write config json: {e}"))?;

        for (name, seq) in self.sequences.iter() {
            let csv_path = dir.join(format!("{name}.csv"));
            seq.to_csv(&csv_path)?;
        }

        Ok(())
    }

    /// Render this sequence machine's state graph as Graphviz DOT text.
    pub fn graphviz_dot(&self) -> String {
        fn dot_quote_id(id: &str) -> String {
            let escaped = id.replace('\\', "\\\\").replace('"', "\\\"");
            format!("\"{escaped}\"")
        }

        fn dot_escape_label(label: &str) -> String {
            let mut escaped = String::new();
            for ch in label.chars() {
                match ch {
                    '\\' => escaped.push_str("\\\\"),
                    '"' => escaped.push_str("\\\""),
                    '\n' => escaped.push_str("\\n"),
                    _ => escaped.push(ch),
                }
            }
            escaped
        }

        fn sequence_node_id(name: &str) -> String {
            format!("seq::{name}")
        }

        fn criterion_label(criterion: &Transition) -> String {
            fn thresh_expr(lhs: &str, op: &ThreshOp, rhs: &str) -> String {
                match op {
                    ThreshOp::Gt { by } => {
                        if *by == 0.0 {
                            format!("{lhs} > {rhs}")
                        } else {
                            format!("{lhs} > {rhs} + {by}")
                        }
                    }
                    ThreshOp::Lt { by } => {
                        if *by == 0.0 {
                            format!("{lhs} < {rhs}")
                        } else {
                            format!("{lhs} < {rhs} - {by}")
                        }
                    }
                    ThreshOp::Approx { atol } => {
                        format!("|{lhs} - {rhs}| < {atol}")
                    }
                }
            }

            match criterion {
                Transition::ConstantThresh(channel, op, threshold) => {
                    thresh_expr(channel, op, &threshold.to_string())
                }
                Transition::ChannelThresh(channel, op, threshold_channel) => {
                    thresh_expr(channel, op, threshold_channel)
                }
                Transition::LookupThresh(channel, op, lookup) => {
                    let expr = thresh_expr(channel, op, "lookup(t)");
                    let t_start = lookup.time_s.first().copied().unwrap_or(f64::NAN);
                    let t_end = lookup.time_s.last().copied().unwrap_or(f64::NAN);
                    format!(
                        "{expr}\nlookup: {} [{:.3e}, {:.3e}]",
                        lookup.method.to_str(),
                        t_start,
                        t_end
                    )
                }
            }
        }

        let mut state_names: BTreeSet<String> = BTreeSet::new();
        state_names.extend(self.sequences.keys().cloned());
        state_names.extend(self.cfg.timeouts.keys().cloned());
        state_names.extend(self.cfg.transitions.keys().cloned());
        for timeout in self.cfg.timeouts.values() {
            if let Timeout::Transition(target) = timeout {
                state_names.insert(target.clone());
            }
        }
        for targets in self.cfg.transitions.values() {
            state_names.extend(targets.keys().cloned());
        }
        if !self.cfg.entry.is_empty() {
            state_names.insert(self.cfg.entry.clone());
        }

        let mut dot = String::new();
        dot.push_str("digraph sequence_machine {\n");
        dot.push_str("  rankdir=LR;\n");
        dot.push_str("  graph [concentrate=true, ranksep=\"1.2 equally\"];\n");
        dot.push_str("  node [shape=box, fontname=\"monospace\", margin=\"0.5,0.10\"];\n");
        dot.push_str("  edge [fontname=\"monospace\"];\n\n");

        for state in &state_names {
            let node_id = sequence_node_id(state);
            if let Some(seq) = self.sequences.get(state) {
                let label = format!(
                    "{state}\n[{:.3e}, {:.3e}] s",
                    seq.get_start_time_s(),
                    seq.get_end_time_s()
                );
                let _ = writeln!(
                    dot,
                    "  {} [label=\"{}\"];",
                    dot_quote_id(&node_id),
                    dot_escape_label(&label)
                );
            } else {
                let label = format!("{state}\n(missing sequence data)");
                let _ = writeln!(
                    dot,
                    "  {} [style=dashed, color=\"red\", label=\"{}\"];",
                    dot_quote_id(&node_id),
                    dot_escape_label(&label)
                );
            }
        }

        if !self.cfg.entry.is_empty() {
            let entry_node_id = "__entry";
            let _ = writeln!(
                dot,
                "  {} [shape=point, width=0.15, label=\"\"];",
                dot_quote_id(entry_node_id)
            );
            let _ = writeln!(
                dot,
                "  {} -> {} [label=\"entry\"];",
                dot_quote_id(entry_node_id),
                dot_quote_id(&sequence_node_id(&self.cfg.entry))
            );
        }
        dot.push('\n');

        for (source, timeout) in self.cfg.timeouts.iter() {
            let source_id = sequence_node_id(source);
            match timeout {
                Timeout::Transition(target) => {
                    let _ = writeln!(
                        dot,
                        "  {} -> {} [style=dashed, label=\"timeout\"];",
                        dot_quote_id(&source_id),
                        dot_quote_id(&sequence_node_id(target))
                    );
                }
                Timeout::Loop => {
                    let _ = writeln!(
                        dot,
                        "  {} -> {} [style=dashed, label=\"timeout loop\"];",
                        dot_quote_id(&source_id),
                        dot_quote_id(&source_id)
                    );
                }
            }
        }

        for (source, targets) in self.cfg.transitions.iter() {
            for (target, criteria) in targets.iter() {
                for criterion in criteria {
                    let _ = writeln!(
                        dot,
                        "  {} -> {} [label=\"{}\"];",
                        dot_quote_id(&sequence_node_id(source)),
                        dot_quote_id(&sequence_node_id(target)),
                        dot_escape_label(&criterion_label(criterion))
                    );
                }
            }
        }

        dot.push_str("}\n");
        dot
    }
}

#[typetag::serde]
impl Calc for SequenceMachine {
    /// Reset internal sequence and register calc tape indices
    fn init(
        &mut self,
        ctx: ControllerCtx,
        input_indices: Vec<usize>,
        output_range: Range<usize>,
    ) -> Result<(), String> {
        // Reload from folder, if linked
        if let Some(rel_path) = &self.cfg.link_folder {
            let folder = ctx.op_dir.join(rel_path);
            *self = *Self::load_folder(&folder)
                .map_err(|e| format!("Failed to load sequence machine from linked folder: {e}"))?;
        }

        // Reset execution sequence
        self.terminate()?;

        // Set per-run config
        self.execution_state.input_indices = input_indices;
        self.execution_state.output_range = output_range;
        self.execution_state.dt_s = ctx.dt_ns as f64 / 1e9;

        // Permute order of each sequence's lookups to match the entrypoint
        let entry_order: Vec<String> = self.current_sequence().data.keys().cloned().collect();
        for s in self.sequences.values_mut() {
            s.permute(&entry_order);
        }

        // Set up map from input names to tape indices to support
        // transition checks
        self.execution_state.input_index_map = BTreeMap::new();
        for (i, name) in self
            .execution_state
            .input_indices
            .iter()
            .cloned()
            .zip(self.get_input_names().iter())
        {
            self.execution_state.input_index_map.insert(name.clone(), i);
        }

        // Make sure lookup tables are usable, transitions refer to real sequences, etc
        self.validate()
    }

    fn terminate(&mut self) -> Result<(), String> {
        self.execution_state.input_indices.clear();
        self.execution_state.output_range = usize::MAX..usize::MAX;
        let start_time = self
            .sequences
            .get(&self.cfg.entry)
            .ok_or_else(|| "Missing sequence".to_string())?
            .get_start_time_s();
        self.execution_state.sequence_time_s = start_time;
        self.execution_state.current_sequence = self.cfg.entry.clone();
        Ok(())
    }

    fn eval(&mut self, tape: &mut [f64]) -> Result<(), String> {
        // Increment sequence time
        self.execution_state.sequence_time_s += self.execution_state.dt_s;
        // Transition to the next sequence if needed, which may reset sequence time
        self.check_transitions(self.execution_state.sequence_time_s, tape)?;

        // Update output values based on the current sequence
        self.current_sequence().eval(
            self.execution_state.sequence_time_s,
            self.execution_state.output_range.clone(),
            tape,
        );
        Ok(())
    }

    /// Map from input field names (like `v`, without prefix) to the sequence name
    /// that the input should draw from (like `peripheral_0.output_1`, with prefix)
    fn get_input_map(&self) -> BTreeMap<CalcInputName, FieldName> {
        let mut map = BTreeMap::new();

        for transitions in self.cfg.transitions.values() {
            for criteria in transitions.values() {
                for criterion in criteria {
                    let names = criterion.get_input_names();
                    for name in names {
                        map.insert(name.clone(), name);
                    }
                }
            }
        }

        map
    }

    /// Change a value in the input map
    fn update_input_map(&mut self, _field: &str, _source: &str) -> Result<(), String> {
        Err(
            "SequenceMachine input map is derived from sequence transition criterion dependencies"
                .to_string(),
        )
    }

    /// Inputs are the sum of all inputs required by any sequence
    fn get_input_names(&self) -> Vec<CalcInputName> {
        self.get_input_map().keys().cloned().collect()
    }

    /// All sequences have the same outputs
    fn get_output_names(&self) -> Vec<CalcOutputName> {
        let mut output_names = vec!["sequence_time_s".to_owned()];
        self.entry_sequence()
            .data
            .keys()
            .cloned()
            .for_each(|n| output_names.push(n));
        output_names
    }

    /// `sequence_time_s` is in seconds; user-defined data channels have no declared unit.
    fn get_output_units(&self) -> Vec<Option<String>> {
        let n_data = self.entry_sequence().data.len();
        let mut units = vec![Some("s".to_owned())];
        units.extend(std::iter::repeat_n(None, n_data));
        units
    }

    /// Get flag for whether to save outputs
    fn get_save_outputs(&self) -> bool {
        self.cfg.save_outputs
    }

    /// Set flag for whether to save outputs
    fn set_save_outputs(&mut self, save_outputs: bool) {
        self.cfg.save_outputs = save_outputs;
    }

    /// Get config field values
    fn get_config(&self) -> BTreeMap<String, f64> {
        BTreeMap::<String, f64>::new()
    }

    /// Apply config field values
    #[allow(unused)]
    fn set_config(&mut self, cfg: &BTreeMap<String, f64>) -> Result<(), String> {
        Err("No settable config fields".to_string())
    }
}

#[cfg(feature = "python")]
#[pymethods]
impl SequenceMachine {
    #[new]
    fn py_new(entry: String) -> Self {
        let cfg = MachineCfg {
            save_outputs: true,
            entry,
            ..Default::default()
        };

        Self {
            cfg,
            sequences: BTreeMap::new(),
            execution_state: ExecutionState::default(),
        }
    }

    /// Serialize to typetagged JSON so Python can pass into trait handoff
    fn to_json(&self) -> PyResult<String> {
        let payload: &dyn Calc = self;
        serde_json::to_string(payload)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
    }

    /// Deserialize from typetagged JSON
    #[staticmethod]
    fn from_json(s: &str) -> PyResult<Self> {
        serde_json::from_str::<Self>(s)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
    }

    #[staticmethod]
    #[pyo3(name = "load_folder")]
    fn py_load_folder(path: &str) -> PyResult<Self> {
        let path = Path::new(path);
        Self::load_folder(&path)
            .map(|machine| *machine)
            .map_err(pyo3::exceptions::PyValueError::new_err)
    }

    #[pyo3(name = "save_folder")]
    fn py_save_folder(&self, path: &str) -> PyResult<()> {
        let path = Path::new(path);
        Self::save_folder(self, &path).map_err(pyo3::exceptions::PyValueError::new_err)
    }

    fn get_entry(&self) -> PyResult<String> {
        Ok(self.cfg.entry.clone())
    }

    fn set_entry(&mut self, entry: String) -> PyResult<()> {
        if !self.sequences.is_empty() && !self.sequences.contains_key(&entry) {
            return Err(pyo3::exceptions::PyKeyError::new_err(format!(
                "Unknown entry sequence: {entry}"
            )));
        }
        self.cfg.entry = entry;
        Ok(())
    }

    fn get_link_folder(&self) -> PyResult<Option<String>> {
        Ok(self.cfg.link_folder.clone())
    }

    fn set_link_folder(&mut self, link_folder: Option<String>) -> PyResult<()> {
        self.cfg.link_folder = link_folder;
        Ok(())
    }

    /// Render this sequence machine's state graph as Graphviz DOT text.
    #[pyo3(name = "graphviz_dot")]
    fn py_graphviz_dot(&self) -> PyResult<String> {
        Ok(self.graphviz_dot())
    }

    fn get_timeout(&self, sequence: String) -> PyResult<Option<String>> {
        let timeout = self.cfg.timeouts.get(&sequence).ok_or_else(|| {
            pyo3::exceptions::PyKeyError::new_err(format!("Unknown sequence: {sequence}"))
        })?;

        match timeout {
            Timeout::Loop => Ok(None),
            Timeout::Transition(target) => Ok(Some(target.clone())),
        }
    }

    fn set_timeout(&mut self, sequence: String, target: Option<String>) -> PyResult<()> {
        if !self.sequences.contains_key(&sequence) {
            return Err(pyo3::exceptions::PyKeyError::new_err(format!(
                "Unknown sequence: {sequence}"
            )));
        }

        let timeout = match target {
            Some(target_sequence) => {
                if !self.sequences.contains_key(&target_sequence) {
                    return Err(pyo3::exceptions::PyKeyError::new_err(format!(
                        "Unknown target sequence: {target_sequence}"
                    )));
                }
                Timeout::Transition(target_sequence)
            }
            None => Timeout::Loop,
        };

        self.cfg.timeouts.insert(sequence, timeout);
        Ok(())
    }

    fn add_sequence(
        &mut self,
        name: String,
        tables: BTreeMap<String, (Vec<f64>, Vec<f64>, String)>,
        timeout: Option<String>,
    ) -> PyResult<()> {
        // Data is required
        if tables.is_empty() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "Sequence data is empty".to_string(),
            ));
        }

        // Check if this sequence is already defined
        if self.sequences.contains_key(&name) {
            return Err(pyo3::exceptions::PyKeyError::new_err(format!(
                "Sequence already exists: {name}"
            )));
        }

        // Build lookups from data
        let mut data = BTreeMap::new();
        for (name, (time_s, vals, method)) in tables {
            // Parse interpolation method
            let method = InterpMethod::try_parse(&method).map_err(|e| {
                pyo3::exceptions::PyValueError::new_err(format!(
                    "Output `{name}` has invalid interp method: {e}"
                ))
            })?;

            // Build interpolator
            let lookup = SequenceLookup::new(method, time_s, vals).map_err(|e| {
                pyo3::exceptions::PyValueError::new_err(format!(
                    "Output `{name}` has invalid lookup data: {e}"
                ))
            })?;
            data.insert(name, lookup);
        }

        if let Some(existing) = self.sequences.values().next() {
            // Make sure all the sequences have the same columns
            let expected: HashSet<String> = existing.data.keys().cloned().collect();
            let provided: HashSet<String> = data.keys().cloned().collect();

            // If the columns don't match, give an informative error
            if expected != provided {
                let mut missing: Vec<String> = expected.difference(&provided).cloned().collect();
                let mut extra: Vec<String> = provided.difference(&expected).cloned().collect();

                // List discrepancies in a consistent order to avoid confusing messages during troubleshooting
                missing.sort();
                extra.sort();

                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "Sequence outputs must match existing sequences. Missing: {missing:?}. Extra: {extra:?}"
                )));
            }
        }

        // Build and validate the combined sequence
        let sequence = Sequence { data };
        sequence.validate().map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("Invalid Sequence: {e:?}"))
        })?;

        // Add the sequence and its timeout setting to the machine
        self.sequences.insert(name.clone(), sequence);
        let timeout = match timeout {
            Some(target_state) => Timeout::Transition(target_state),
            None => Timeout::Loop,
        };
        self.cfg.timeouts.insert(name.clone(), timeout);
        self.cfg.transitions.entry(name).or_default();

        Ok(())
    }

    /// Add a constant threshold transition for a sequence.
    fn add_constant_thresh_transition(
        &mut self,
        source_target: (String, String),
        channel: String,
        op: (&str, f64),
        threshold: f64,
    ) -> PyResult<()> {
        // Unpack and parse
        let (source_sequence, target_sequence) = source_target;
        let op = ThreshOp::try_parse(op).map_err(pyo3::exceptions::PyValueError::new_err)?;

        // Add to machine
        let transition = Transition::ConstantThresh(channel, op, threshold);
        self.add_transition(source_sequence, target_sequence, transition)
            .map_err(pyo3::exceptions::PyValueError::new_err)
    }

    /// Add a channel threshold transition for a sequence.
    fn add_channel_thresh_transition(
        &mut self,
        source_target: (String, String),
        channel: String,
        op: (&str, f64),
        threshold_channel: String,
    ) -> PyResult<()> {
        // Unpack and parse
        let (source_sequence, target_sequence) = source_target;
        let op = ThreshOp::try_parse(op).map_err(pyo3::exceptions::PyValueError::new_err)?;

        // Add to machine
        let transition = Transition::ChannelThresh(channel, op, threshold_channel);
        self.add_transition(source_sequence, target_sequence, transition)
            .map_err(pyo3::exceptions::PyValueError::new_err)
    }

    /// Add a lookup threshold transition for a sequence.
    fn add_lookup_thresh_transition(
        &mut self,
        source_target: (String, String),
        channel: String,
        op: (&str, f64),
        threshold_lookup: (Vec<f64>, Vec<f64>, &str),
    ) -> PyResult<()> {
        // Unpack and parse
        let (source_sequence, target_sequence) = source_target;
        let (time_s, vals, method) = threshold_lookup;
        let op = ThreshOp::try_parse(op).map_err(pyo3::exceptions::PyValueError::new_err)?;
        let method = InterpMethod::try_parse(method).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!(
                "Lookup has invalid interp method: {e}"
            ))
        })?;

        // Build lookup interpolator
        let lookup = SequenceLookup::new(method, time_s, vals).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("Lookup has invalid data: {e}"))
        })?;

        // Add to machine
        let transition = Transition::LookupThresh(channel, op, lookup);
        self.add_transition(source_sequence, target_sequence, transition)
            .map_err(pyo3::exceptions::PyValueError::new_err)
    }
}

#[cfg(test)]
mod tests {
    use super::SequenceMachine;
    use std::path::PathBuf;

    /// Check that we can save and load the human-readable folder format
    /// without losing or corrupting information
    #[test]
    fn roundtrip_sequence_machine_folder() {
        let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        let src_dir = root.join("examples").join("machine");
        let tmp_dir = std::env::temp_dir().join("deimos_sequence_roundtrip");

        let _ = std::fs::remove_dir_all(&tmp_dir);
        std::fs::create_dir_all(&tmp_dir).unwrap();

        let original = *SequenceMachine::load_folder(&src_dir).unwrap();
        let original_json = serde_json::to_string_pretty(&original).unwrap();

        original.save_folder(&tmp_dir).unwrap();
        let roundtrip = *SequenceMachine::load_folder(&tmp_dir).unwrap();
        let roundtrip_json = serde_json::to_string_pretty(&roundtrip).unwrap();

        assert_eq!(original_json, roundtrip_json);
    }

    #[test]
    fn graphviz_dot_contains_entry_and_transitions() {
        let root = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        let src_dir = root.join("examples").join("machine");
        let machine = *SequenceMachine::load_folder(&src_dir).unwrap();

        let dot = machine.graphviz_dot();

        assert!(dot.contains("digraph sequence_machine"));
        assert!(dot.contains("\"__entry\" -> \"seq::low\" [label=\"entry\"];"));
        assert!(dot.contains("\"seq::low\" -> \"seq::high\""));
        assert!(dot.contains("\"seq::high\" -> \"seq::low\" [style=dashed, label=\"timeout\"];"));
    }
}