earthsci-ast 0.6.0

//! Reaction system analysis and ODE generation (esm-spec §7.4).
//!
//! Two public surfaces: the stoichiometric-matrix builders
//! ([`stoichiometric_matrix`] and friends) and mass-action ODE derivation
//! ([`derive_reaction_odes`], §4.7.5), which turns a [`ReactionSystem`]'s
//! reactions into per-species `D(sp, t) = Σ ±coeff·rate·Π substrates`
//! equations mirroring the Julia (Catalyst) and Python references.

use crate::{
    Equation, Expr, ExpressionNode, InlineValue, Model, ModelVariable, ReactionSystem, Species,
    VariableType,
};
use indexmap::IndexMap;
use std::collections::HashMap;
use thiserror::Error;

/// Error type for ODE derivation operations
#[derive(Error, Debug)]
pub enum DeriveError {
    /// Invalid stoichiometry
    #[error("Invalid stoichiometry: {0}")]
    InvalidStoichiometry(String),
}

/// Lower a reaction network to an ODE equation list.
///
/// Produces one `D(species, t) = Σ(net_stoichiometry · rate)` equation per species
/// per spec §4.7.5 step 1. Reusable from both [`derive_odes`] and
/// [`crate::flatten::flatten`] so the reaction-to-equation core lives in one place.
///
/// Each reaction's base rate is enhanced with mass-action concentration factors
/// per spec §7.4 (the `rate` field is the coefficient; the runner always
/// multiplies by the substrate product). Net stoichiometry combines substrate
/// and product contributions for each species.
pub fn lower_reactions_to_equations(
    reactions: &[crate::Reaction],
    species: &IndexMap<String, Species>,
) -> Result<Vec<Equation>, DeriveError> {
    // Validate system has species if there are reactions
    if species.is_empty() && !reactions.is_empty() {
        return Err(DeriveError::InvalidStoichiometry(
            "Reaction system has reactions but no species defined".to_string(),
        ));
    }

    // Validate reactions and their stoichiometry. Coefficients are f64
    // (fractional stoichiometries are allowed); non-negativity is enforced
    // upstream by the schema (`exclusiveMinimum: 0`) and the parse-level
    // stoichiometry check, so only species references are checked here.
    for (reaction_idx, reaction) in reactions.iter().enumerate() {
        for substrate in reaction.substrates.iter().flatten() {
            if !species.contains_key(&substrate.species) {
                return Err(DeriveError::InvalidStoichiometry(format!(
                    "Unknown substrate species '{}' in reaction {}",
                    substrate.species, reaction_idx
                )));
            }
        }

        for product in reaction.products.iter().flatten() {
            if !species.contains_key(&product.species) {
                return Err(DeriveError::InvalidStoichiometry(format!(
                    "Unknown product species '{}' in reaction {}",
                    product.species, reaction_idx
                )));
            }
        }

        let no_substrates = reaction.substrates.as_ref().is_none_or(|v| v.is_empty());
        let no_products = reaction.products.as_ref().is_none_or(|v| v.is_empty());
        if no_substrates && no_products {
            return Err(DeriveError::InvalidStoichiometry(format!(
                "Reaction {reaction_idx} has no substrates or products"
            )));
        }
    }

    // Each species' terms, indexed by its declaration position and filled
    // reaction by reaction so they come out in reaction order
    // (esm-libraries-spec §4.6.2 rule 4).
    let mut species_terms: Vec<Vec<Expr>> = vec![Vec::new(); species.len()];

    for reaction in reactions {
        // v0.2.x allows fractional stoichiometries; accumulate in f64 so
        // products like `0.87 CH2O` survive the ODE lowering unchanged.
        // The net coefficient is `P - S` (esm-libraries-spec §4.6.2): each
        // species' product and substrate coefficients are summed in listed
        // order and differenced once. Only the species this reaction names are
        // visited, as `(index, produced, consumed)` in first-mention order.
        let mut sums: Vec<(usize, f64, f64)> = Vec::new();
        let sides = [
            (false, reaction.substrates.iter().flatten()),
            (true, reaction.products.iter().flatten()),
        ];
        for (is_product, entries) in sides {
            for entry in entries {
                // Validated above: every entry names a declared species.
                let Some(idx) = species.get_index_of(&entry.species) else {
                    continue;
                };
                let slot = match sums.iter().position(|(i, _, _)| *i == idx) {
                    Some(pos) => &mut sums[pos],
                    None => {
                        sums.push((idx, 0.0, 0.0));
                        sums.last_mut().unwrap()
                    }
                };
                if is_product {
                    slot.1 += entry.coefficient;
                } else {
                    slot.2 += entry.coefficient;
                }
            }
        }

        // The mass-action rate law, built once per reaction and cloned into
        // each species' term.
        let mut rate_law: Option<Expr> = None;
        for (idx, produced, consumed) in sums {
            // Reservoir species (`constant: true`, §7.4): held fixed, so it
            // gets NO dX/dt equation. Its concentration is still a mass-action
            // factor in the OTHER species' rate laws — the rate law reads it
            // from the reaction's substrate list regardless of this skip.
            // Mirrors the Julia reference (flatten.jl
            // `lower_reactions_to_equations`).
            if species[idx].constant == Some(true) {
                continue;
            }
            let net_stoichiometry = produced - consumed;
            if net_stoichiometry == 0.0 {
                continue;
            }
            let enhanced_rate = rate_law
                .get_or_insert_with(|| {
                    enhance_rate_with_mass_action(
                        &reaction.rate,
                        reaction.substrates.as_deref().unwrap_or(&[]),
                    )
                })
                .clone();

            // The canonical term (esm-libraries-spec §4.6.2): the bare
            // rate law for a net coefficient of exactly 1, otherwise the
            // two-operand product `coefficient * rate`, coefficient first,
            // an INTEGER literal when integral. A loss of 1 is
            // `-1 * rate`; `-1.0` would serialize as a float.
            species_terms[idx].push(if net_stoichiometry == 1.0 {
                enhanced_rate
            } else {
                Expr::operator(ExpressionNode {
                    op: "*".to_string(),
                    args: vec![stoich_literal(net_stoichiometry), enhanced_rate],
                    wrt: None,
                    dim: None,
                    ..Default::default()
                })
            });
        }
    }

    let mut equations = Vec::with_capacity(species.len());

    // DOCUMENT ORDER (esm-libraries-spec §4.7.5 step 4, "Ordering"): the
    // derived `D(species, t)` equations follow the order the reaction system
    // DECLARES its species. Sorting them — what this used to do — reorders the
    // flattened equation list relative to `state_variables`, which the shared
    // flatten corpus pins per equation index.
    for (sp_name, rate_terms) in species.keys().zip(species_terms) {
        // A species with NO net contribution from any reaction gets NO
        // equation, rather than a vacuous `D(X, t) = 0`. Emitting the zero
        // equation makes such a species an ODE state that never moves, and
        // makes `equation_count` incomparable across bindings; omitting it
        // leaves the species an unknown that no equation names, which
        // esm-spec §6.3.1's partition-by-complement classifies ALGEBRAIC —
        // the oracle's answer, and the one Rust's own `classification`
        // module already reaches for such an unknown.
        if rate_terms.is_empty() {
            continue;
        }

        let rhs = if rate_terms.len() == 1 {
            rate_terms.into_iter().next().unwrap()
        } else {
            Expr::operator(ExpressionNode {
                op: "+".to_string(),
                args: rate_terms,
                wrt: None,
                dim: None,
                ..Default::default()
            })
        };

        let lhs = Expr::operator(ExpressionNode {
            op: "D".to_string(),
            args: vec![Expr::Variable(sp_name.clone())],
            wrt: Some("t".to_string()),
            dim: None,
            ..Default::default()
        });

        equations.push(Equation {
            comment: None,
            lhs,
            rhs,
        });
    }

    Ok(equations)
}

/// A stoichiometric coefficient or exponent as the literal the lowered tree
/// carries: [`Expr::Integer`] when the value is integral and representable in
/// `i64` (CONFORMANCE_SPEC §5.5.3.1 rule 3, esm-libraries-spec §4.6.2 rule 5),
/// [`Expr::Number`] otherwise. `i64::MAX as f64` rounds up to 2^63, one past
/// `i64::MAX`, so the upper bound is exclusive.
fn stoich_literal(value: f64) -> Expr {
    if value.fract() == 0.0 && (i64::MIN as f64..i64::MAX as f64).contains(&value) {
        Expr::Integer(value as i64)
    } else {
        Expr::Number(value)
    }
}

/// Generate ODE model from a reaction system
///
/// Converts a reaction system into an ODE model with species as state variables
/// and reactions contributing to their derivatives using mass action kinetics.
///
/// Mass action kinetics: rate law = k * product(substrates^stoichiometry)
/// Net stoichiometry = products - substrates
/// `d[species]/dt = sum(net_stoichiometry * rate_law)`
///
/// # Arguments
///
/// * `system` - The reaction system to convert
///
/// # Returns
///
/// * `Result<Model, DeriveError>` - ODE model with species as state variables, or error
///
/// # Errors
///
/// Returns `DeriveError` for invalid stoichiometry, missing rate laws, or unit conversion issues.
pub fn derive_odes(system: &ReactionSystem) -> Result<Model, DeriveError> {
    let mut variables = IndexMap::new();

    for (species_name, species) in &system.species {
        // Reservoir species (`constant: true`, §7.4) are held fixed and get no
        // ODE (skipped in `lower_reactions_to_equations`), so they lower to
        // PARAMETERS carrying their declared `default` as the fixed value — the
        // same treatment as the flatten path (`build_reaction_block`) and the
        // Julia reference (reactions.jl `derive_odes`).
        let var_type = if species.constant == Some(true) {
            VariableType::Parameter
        } else {
            VariableType::Unknown
        };
        variables.insert(
            species_name.clone(),
            ModelVariable {
                var_type,
                units: species.units.clone(),
                default: species.default.map(InlineValue::Scalar),
                description: species.description.clone(),
                ..Default::default()
            },
        );
    }

    // Rate PARAMETERS become parameter variables of the derived model. Without
    // this the derived equations reference rate constants (`k1`, `jNO2`, …)
    // that the model does not declare, so the model does not close — every
    // rate is an undefined variable. Julia (reactions.jl "Add parameter
    // variables for all parameters"), Python and TypeScript ("Copy
    // parameters") all do it; this binding alone did not. Order matches
    // Python and TypeScript: species in declaration order, then parameters in
    // declaration order.
    //
    // The §6.3 value model travels with them. `update` above all: it is the
    // only channel binding a parameter to a data source (esm-spec §5.4), so
    // dropping it here would lower a data-driven parameter to a constant.
    for (param_name, param) in &system.parameters {
        variables.insert(
            param_name.clone(),
            ModelVariable {
                var_type: VariableType::Parameter,
                units: param.units.clone(),
                default: param.default.map(InlineValue::Scalar),
                description: param.description.clone(),
                default_units: param.default_units.clone(),
                shape: param.shape.clone(),
                distribution: param.distribution.clone(),
                update: param.update.clone(),
                ..Default::default()
            },
        );
    }

    let equations = lower_reactions_to_equations(&system.reactions, &system.species)?;

    Ok(Model {
        reference: system.reference.clone(),
        variables,
        equations,
        discrete_events: system.discrete_events.clone(),
        continuous_events: system.continuous_events.clone(),
        ..Default::default()
    })
}

/// Enhance base rate law with mass action kinetics
///
/// Per `esm-spec.md` §7.4 the `rate` field is the rate COEFFICIENT.
/// The full rate law is always `k * product(substrates^stoichiometry)` — the
/// runner unconditionally multiplies the coefficient by the substrate product.
/// Source reactions (no substrates) return the coefficient unchanged.
///
/// Infallible — every path yields an [`Expr`] (the mass-action lowering cannot
/// fail once the substrates are known), so it returns a bare `Expr` rather than
/// a `Result`.
fn enhance_rate_with_mass_action(rate: &Expr, substrates: &[crate::StoichiometricEntry]) -> Expr {
    // If no substrates (source reaction), return rate as-is
    if substrates.is_empty() {
        return rate.clone();
    }

    // Enhance with mass action kinetics (spec §7.4): `∏ [Sᵢ]^nᵢ`. Mirrors
    // the Julia reference (reactions.jl) exactly — a coefficient of 1 is the
    // bare species reference, an integral coefficient lowers to `^` with an
    // Integer exponent node, and a fractional coefficient (e.g. 1.5) to `^`
    // with a Number exponent. (An earlier Rust-only behavior unrolled small
    // integer coefficients into repeated multiplication `[A]·[A]·…`, which
    // diverged from both sibling bindings' ASTs and built enormous nodes for
    // large coefficients.)
    let mut concentration_factors = Vec::new();

    for substrate in substrates {
        let coeff = substrate.coefficient;
        let species_var = Expr::Variable(substrate.species.clone());

        if coeff == 1.0 {
            concentration_factors.push(species_var);
        } else {
            let exponent = stoich_literal(coeff);
            concentration_factors.push(Expr::operator(ExpressionNode {
                op: "^".to_string(),
                args: vec![species_var, exponent],
                wrt: None,
                dim: None,
                ..Default::default()
            }));
        }
    }

    // Combine rate coefficient with concentration factors
    let mut all_factors = vec![rate.clone()];
    all_factors.extend(concentration_factors);

    if all_factors.len() == 1 {
        all_factors.into_iter().next().unwrap()
    } else {
        Expr::operator(ExpressionNode {
            op: "*".to_string(),
            args: all_factors,
            wrt: None,
            dim: None,
            ..Default::default()
        })
    }
}

/// Generate stoichiometric matrix from a reaction system
///
/// Creates a matrix where rows represent species and columns represent reactions.
/// `matrix[i][j]` is the stoichiometric coefficient of species `i` in reaction `j`.
/// Negative values indicate reactants, positive values indicate products.
///
/// # Arguments
///
/// * `system` - The reaction system to analyze
///
/// # Returns
///
/// * `Vec<Vec<f64>>` - Matrix with species as rows and reactions as columns
pub fn stoichiometric_matrix(system: &ReactionSystem) -> Vec<Vec<f64>> {
    let num_species = system.species.len();
    let num_reactions = system.reactions.len();

    // Initialize matrix with zeros
    let mut matrix = vec![vec![0.0f64; num_reactions]; num_species];

    // ROW ORDER IS DECLARATION ORDER (API_SPEC.md §5.10). `system.species` is an
    // `IndexMap`, so iterating its keys is already deterministic across runs AND
    // already the order the document declares. This used to `.sort()` the names
    // first, justified as "so indices are reproducible" — but reproducibility is
    // what the container was chosen to give, so the sort bought nothing and cost
    // the agreement with `derive_odes` / `lower_reactions_to_equations`, which
    // never sorted, and with Julia, Python and TypeScript, which never sorted
    // either.
    let species_index: HashMap<String, usize> = system
        .species
        .keys()
        .enumerate()
        .map(|(idx, name)| (name.clone(), idx))
        .collect();

    // Fill in the matrix
    for (reaction_idx, reaction) in system.reactions.iter().enumerate() {
        // Process substrates (negative coefficients)
        for substrate in reaction.substrates.iter().flatten() {
            if let Some(&species_idx) = species_index.get(&substrate.species) {
                matrix[species_idx][reaction_idx] -= substrate.coefficient;
            }
        }

        // Process products (positive coefficients)
        for product in reaction.products.iter().flatten() {
            if let Some(&species_idx) = species_index.get(&product.species) {
                matrix[species_idx][reaction_idx] += product.coefficient;
            }
        }
    }

    matrix
}

/// Generate a stoichiometric matrix from a reaction system using parallel computation
///
/// This function provides parallel computation for generating stoichiometric matrices,
/// which can significantly improve performance for large reaction systems.
///
/// # Arguments
///
/// * `system` - The reaction system to analyze
///
/// # Returns
///
/// * `Result<Vec<Vec<f64>>, crate::performance::PerformanceError>` - Matrix with species as rows and reactions as columns
///
/// # Features
///
/// This function requires the `parallel` feature to be enabled.
///
/// # Example
///
/// ```rust
/// # #[cfg(feature = "parallel")]
/// # {
/// use earthsci_ast::{ReactionSystem, stoichiometric_matrix_parallel};
///
/// // Assuming you have a reaction system
/// let reaction_system = ReactionSystem::default();
///
/// let matrix = stoichiometric_matrix_parallel(&reaction_system).unwrap();
/// # }
/// ```
#[cfg(feature = "parallel")]
pub fn stoichiometric_matrix_parallel(
    system: &ReactionSystem,
) -> Result<Vec<Vec<f64>>, crate::performance::PerformanceError> {
    use crate::performance::ParallelEvaluator;

    let evaluator = ParallelEvaluator::new(None)?;
    evaluator.compute_stoichiometric_matrix_parallel(system)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Reaction, Species, StoichiometricEntry, VariableType};
    use indexmap::IndexMap;

    #[test]
    fn test_stoich_literal_int64_bounds() {
        // 2^63 is integral but one past `i64::MAX`: a float, not a saturated
        // integer (esm-libraries-spec §4.6.2 rule 5).
        let two_63 = 2.0_f64.powi(63);
        assert!(matches!(stoich_literal(two_63), Expr::Number(_)));
        assert!(matches!(stoich_literal(-two_63), Expr::Integer(i64::MIN)));
        assert!(matches!(
            stoich_literal(1e17),
            Expr::Integer(100_000_000_000_000_000)
        ));
        assert!(matches!(stoich_literal(-2.0), Expr::Integer(-2)));
        assert!(matches!(stoich_literal(1.5), Expr::Number(_)));
    }

    fn create_test_species(name: &str) -> (String, Species) {
        (
            name.to_string(),
            Species {
                default_units: None,
                units: Some("mol/L".to_string()),
                default: Some(0.0),
                description: None,
                constant: None,
            },
        )
    }

    fn create_test_reaction(
        substrates: Vec<(&str, f64)>,
        products: Vec<(&str, f64)>,
        rate: Expr,
    ) -> Reaction {
        Reaction {
            id: None,
            name: None,
            substrates: Some(
                substrates
                    .into_iter()
                    .map(|(species, coeff)| StoichiometricEntry {
                        species: species.to_string(),
                        coefficient: coeff,
                    })
                    .collect(),
            ),
            products: Some(
                products
                    .into_iter()
                    .map(|(species, coeff)| StoichiometricEntry {
                        species: species.to_string(),
                        coefficient: coeff,
                    })
                    .collect(),
            ),
            rate,
            reference: None,
        }
    }

    /// The derived model must DECLARE the rate parameters its equations
    /// reference. Julia, Python and TypeScript all copy `parameters` into the
    /// derived `variables`; this binding did not, so every rate constant came
    /// out an undefined variable.
    #[test]
    fn derive_odes_declares_rate_parameters() {
        let mut species = IndexMap::new();
        species.insert(
            "A".to_string(),
            Species {
                default_units: None,
                default: Some(1.0),
                units: Some("molec/cm3".to_string()),
                description: None,
                constant: None,
            },
        );

        let mut parameters = IndexMap::new();
        parameters.insert(
            "k1".to_string(),
            crate::types::Parameter {
                default_units: None,
                shape: None,
                distribution: None,
                update: None,
                default: Some(1.8e-12),
                units: Some("cm3/molec/s".to_string()),
                description: Some("Rate constant".to_string()),
            },
        );

        let system = ReactionSystem {
            species,
            parameters,
            reactions: vec![Reaction {
                id: Some("R1".to_string()),
                name: None,
                substrates: Some(vec![StoichiometricEntry {
                    species: "A".to_string(),
                    coefficient: 1.0,
                }]),
                products: None,
                rate: Expr::Variable("k1".to_string()),
                reference: None,
            }],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("derives");
        let k1 = model.variables.get("k1").expect("k1 is declared");
        assert!(matches!(k1.var_type, VariableType::Parameter));
        assert_eq!(k1.default_scalar(), Some(1.8e-12));
        assert_eq!(k1.units.as_deref(), Some("cm3/molec/s"));

        // Species first in declaration order, then parameters -- the order
        // Python and TypeScript produce.
        let names: Vec<&str> = model.variables.keys().map(String::as_str).collect();
        assert_eq!(names, vec!["A", "k1"]);
    }

    #[test]
    fn test_derive_odes_simple() {
        let system = ReactionSystem {
            species: [create_test_species("A"), create_test_species("B")]
                .into_iter()
                .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // A -> B with rate k1 * A
                create_test_reaction(
                    vec![("A", 1.0)],
                    vec![("B", 1.0)],
                    Expr::operator(ExpressionNode {
                        op: "*".to_string(),
                        args: vec![
                            Expr::Variable("k1".to_string()),
                            Expr::Variable("A".to_string()),
                        ],
                        wrt: None,
                        dim: None,
                        ..Default::default()
                    }),
                ),
            ],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should derive ODEs successfully");

        assert_eq!(model.variables.len(), 2);
        assert!(model.variables.contains_key("A"));
        assert!(model.variables.contains_key("B"));

        assert_eq!(model.equations.len(), 2);

        // Both species should have ODE equations
        let var_names: Vec<String> = model
            .equations
            .iter()
            .map(|eq| match &eq.lhs {
                Expr::Operator(node) if node.op == "D" => match &node.args[0] {
                    Expr::Variable(name) => name.clone(),
                    _ => "unknown".to_string(),
                },
                _ => "unknown".to_string(),
            })
            .collect();

        assert!(var_names.contains(&"A".to_string()));
        assert!(var_names.contains(&"B".to_string()));
    }

    #[test]
    fn test_derive_odes_reservoir_constant_species() {
        // esm-spec §7.4: a `constant: true` reservoir species that is a
        // stoichiometric SUBSTRATE (R + A -> B) is held fixed — it lowers to a
        // PARAMETER (not a state), gets NO dR/dt equation, yet still appears as
        // a mass-action concentration factor in the other species' rate laws.
        let mut species: IndexMap<String, Species> =
            [create_test_species("A"), create_test_species("B")]
                .into_iter()
                .collect();
        species.insert(
            "R".to_string(),
            Species {
                default_units: None,
                units: Some("1".to_string()),
                default: Some(2.0),
                description: None,
                constant: Some(true),
            },
        );
        let system = ReactionSystem {
            species,
            parameters: IndexMap::new(),
            reactions: vec![create_test_reaction(
                vec![("R", 1.0), ("A", 1.0)],
                vec![("B", 1.0)],
                Expr::Variable("k".to_string()),
            )],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should derive ODEs successfully");

        // Reservoir R becomes a Parameter carrying its default; A and B stay State.
        assert!(
            matches!(model.variables["R"].var_type, VariableType::Parameter),
            "reservoir R must lower to a parameter"
        );
        assert_eq!(model.variables["R"].default_scalar(), Some(2.0));
        assert!(matches!(
            model.variables["A"].var_type,
            VariableType::Unknown
        ));
        assert!(matches!(
            model.variables["B"].var_type,
            VariableType::Unknown
        ));

        // No dR/dt equation; A and B keep theirs (2 equations total, not 3).
        let eq_species: Vec<String> = model
            .equations
            .iter()
            .filter_map(|eq| match &eq.lhs {
                Expr::Operator(node) if node.op == "D" => match &node.args[0] {
                    Expr::Variable(name) => Some(name.clone()),
                    _ => None,
                },
                _ => None,
            })
            .collect();
        assert!(
            !eq_species.contains(&"R".to_string()),
            "reservoir R must get no ODE, got equations for {eq_species:?}"
        );
        assert!(eq_species.contains(&"A".to_string()));
        assert!(eq_species.contains(&"B".to_string()));
        assert_eq!(model.equations.len(), 2);

        // R is still a mass-action factor in dA/dt (held-fixed reservoir still
        // contributes its concentration to every rate law it appears in).
        let a_eq = model
            .equations
            .iter()
            .find(|eq| match &eq.lhs {
                Expr::Operator(node) if node.op == "D" => {
                    matches!(&node.args[0], Expr::Variable(name) if name == "A")
                }
                _ => false,
            })
            .expect("Should find A equation");
        let rhs_dbg = format!("{:?}", a_eq.rhs);
        assert!(
            rhs_dbg.contains("\"R\""),
            "reservoir R must remain a mass-action factor in dA/dt: {rhs_dbg}"
        );
    }

    #[test]
    fn test_stoichiometric_matrix() {
        let system = ReactionSystem {
            species: [
                create_test_species("A"),
                create_test_species("B"),
                create_test_species("C"),
            ]
            .into_iter()
            .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // Reaction 1: A -> B
                create_test_reaction(
                    vec![("A", 1.0)],
                    vec![("B", 1.0)],
                    Expr::Variable("k1".to_string()),
                ),
                // Reaction 2: B -> C
                create_test_reaction(
                    vec![("B", 1.0)],
                    vec![("C", 1.0)],
                    Expr::Variable("k2".to_string()),
                ),
                // Reaction 3: 2A -> C
                create_test_reaction(
                    vec![("A", 2.0)],
                    vec![("C", 1.0)],
                    Expr::Variable("k3".to_string()),
                ),
            ],
            ..Default::default()
        };

        let matrix = stoichiometric_matrix(&system);

        // Should be 3x3 matrix (3 species, 3 reactions)
        assert_eq!(matrix.len(), 3);
        assert_eq!(matrix[0].len(), 3);
        assert_eq!(matrix[1].len(), 3);
        assert_eq!(matrix[2].len(), 3);

        // Check specific values
        // Species A: [-1, 0, -2] (consumed in reactions 1 and 3)
        assert_eq!(matrix[0], vec![-1.0, 0.0, -2.0]);

        // Species B: [1, -1, 0] (produced in reaction 1, consumed in reaction 2)
        assert_eq!(matrix[1], vec![1.0, -1.0, 0.0]);

        // Species C: [0, 1, 1] (produced in reactions 2 and 3)
        assert_eq!(matrix[2], vec![0.0, 1.0, 1.0]);
    }

    #[test]
    fn test_stoichiometric_matrix_empty() {
        let system = ReactionSystem {
            species: indexmap::IndexMap::new(),
            parameters: IndexMap::new(),
            reactions: vec![],
            ..Default::default()
        };

        let matrix = stoichiometric_matrix(&system);
        assert_eq!(matrix.len(), 0);
    }

    #[test]
    fn test_derive_odes_empty_system() {
        let system = ReactionSystem {
            species: indexmap::IndexMap::new(),
            parameters: IndexMap::new(),
            reactions: vec![],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should handle empty system");
        assert_eq!(model.variables.len(), 0);
        assert_eq!(model.equations.len(), 0);
    }

    #[test]
    fn test_derive_odes_unknown_species_error() {
        let system = ReactionSystem {
            species: [create_test_species("A")]
                .into_iter()
                .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![create_test_reaction(
                vec![("B", 1.0)], // B is not defined in species
                vec![("A", 1.0)],
                Expr::Variable("k1".to_string()),
            )],
            ..Default::default()
        };

        let result = derive_odes(&system);
        assert!(result.is_err());
        match result {
            Err(DeriveError::InvalidStoichiometry(msg)) => {
                assert!(msg.contains("Unknown substrate species 'B'"));
            }
            _ => panic!("Expected InvalidStoichiometry error"),
        }
    }

    #[test]
    fn test_derive_odes_mass_action_kinetics() {
        let system = ReactionSystem {
            species: [
                create_test_species("A"),
                create_test_species("B"),
                create_test_species("C"),
            ]
            .into_iter()
            .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // A + B -> C with rate coefficient k1 (should become k1*A*B)
                create_test_reaction(
                    vec![("A", 1.0), ("B", 1.0)],
                    vec![("C", 1.0)],
                    Expr::Variable("k1".to_string()),
                ),
            ],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should derive ODEs successfully");
        assert_eq!(model.variables.len(), 3);
        assert_eq!(model.equations.len(), 3);

        // Check that the rate law includes mass action terms
        // For species C: d[C]/dt = k1 * A * B
        let c_equation = model
            .equations
            .iter()
            .find(|eq| match &eq.lhs {
                Expr::Operator(node) if node.op == "D" => match &node.args[0] {
                    Expr::Variable(name) => name == "C",
                    _ => false,
                },
                _ => false,
            })
            .expect("Should find C equation");

        // The RHS should be a multiplication involving k1, A, and B
        match &c_equation.rhs {
            Expr::Operator(node) if node.op == "*" => {
                assert!(node.args.len() >= 2);
                // Should contain k1, A, and B in some form
            }
            _ => panic!("Expected multiplication for mass action kinetics"),
        }
    }

    #[test]
    fn test_derive_odes_source_reaction() {
        let system = ReactionSystem {
            species: [create_test_species("A")]
                .into_iter()
                .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // Source reaction: -> A with rate k0 (no substrates)
                create_test_reaction(vec![], vec![("A", 1.0)], Expr::Variable("k0".to_string())),
            ],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should handle source reactions");
        assert_eq!(model.variables.len(), 1);
        assert_eq!(model.equations.len(), 1);

        // For species A: d[A]/dt = k0 (no concentration dependence)
        let a_equation = &model.equations[0];
        match &a_equation.rhs {
            Expr::Variable(name) => assert_eq!(name, "k0"),
            _ => panic!("Expected simple rate constant for source reaction"),
        }
    }

    #[test]
    fn test_derive_odes_sink_reaction() {
        let system = ReactionSystem {
            species: [create_test_species("A")]
                .into_iter()
                .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // Sink reaction: A -> with rate k_deg (no products)
                create_test_reaction(
                    vec![("A", 1.0)],
                    vec![],
                    Expr::Variable("k_deg".to_string()),
                ),
            ],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should handle sink reactions");
        assert_eq!(model.variables.len(), 1);
        assert_eq!(model.equations.len(), 1);

        // For species A: d[A]/dt = -k_deg * A
        let a_equation = &model.equations[0];
        match &a_equation.rhs {
            Expr::Operator(node) if node.op == "*" => {
                // Should be [-1, k_deg * A] structure
                assert_eq!(node.args.len(), 2);
                // First arg should be -1, second should be k_deg * A
                match &node.args[0] {
                    // An INTEGER literal (esm-libraries-spec §4.6.2), never
                    // the float `-1.0`.
                    Expr::Integer(n) => assert_eq!(*n, -1),
                    _ => panic!("Expected the integer literal -1 as first argument"),
                }
            }
            _ => panic!(
                "Expected multiplication for sink reaction kinetics, got: {:?}",
                a_equation.rhs
            ),
        }
    }

    #[test]
    fn test_derive_odes_higher_order_reaction() {
        let system = ReactionSystem {
            species: [create_test_species("A"), create_test_species("B")]
                .into_iter()
                .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // 2A -> B with rate k1 (second order in A)
                create_test_reaction(
                    vec![("A", 2.0)],
                    vec![("B", 1.0)],
                    Expr::Variable("k1".to_string()),
                ),
            ],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should handle higher order reactions");
        assert_eq!(model.variables.len(), 2);
        assert_eq!(model.equations.len(), 2);

        // Check that the rate law includes A^2 term (or A*A)
        let b_equation = model
            .equations
            .iter()
            .find(|eq| match &eq.lhs {
                Expr::Operator(node) if node.op == "D" => match &node.args[0] {
                    Expr::Variable(name) => name == "B",
                    _ => false,
                },
                _ => false,
            })
            .expect("Should find B equation");

        // The RHS should involve k1 and A squared
        match &b_equation.rhs {
            Expr::Operator(node) if node.op == "*" => {
                assert!(node.args.len() >= 2);
                // Should contain k1 and either A*A or pow(A, 2)
            }
            _ => panic!("Expected multiplication for higher order kinetics"),
        }
    }

    #[test]
    fn test_derive_odes_reactions_with_no_substrates_and_products() {
        let system = ReactionSystem {
            species: [create_test_species("A")]
                .into_iter()
                .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![create_test_reaction(
                vec![], // No substrates
                vec![], // No products
                Expr::Variable("k1".to_string()),
            )],
            ..Default::default()
        };

        let result = derive_odes(&system);
        assert!(result.is_err());
        match result {
            Err(DeriveError::InvalidStoichiometry(msg)) => {
                assert!(msg.contains("has no substrates or products"));
            }
            _ => panic!("Expected InvalidStoichiometry error"),
        }
    }

    #[test]
    fn test_derive_odes_complex_reaction_network() {
        let system = ReactionSystem {
            species: [
                create_test_species("A"),
                create_test_species("B"),
                create_test_species("C"),
                create_test_species("D"),
            ]
            .into_iter()
            .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // A + B -> C + D (rate k1)
                create_test_reaction(
                    vec![("A", 1.0), ("B", 1.0)],
                    vec![("C", 1.0), ("D", 1.0)],
                    Expr::Variable("k1".to_string()),
                ),
                // C -> A (rate k2)
                create_test_reaction(
                    vec![("C", 1.0)],
                    vec![("A", 1.0)],
                    Expr::Variable("k2".to_string()),
                ),
                // D -> B (rate k3)
                create_test_reaction(
                    vec![("D", 1.0)],
                    vec![("B", 1.0)],
                    Expr::Variable("k3".to_string()),
                ),
            ],
            ..Default::default()
        };

        let model = derive_odes(&system).expect("Should handle complex networks");
        assert_eq!(model.variables.len(), 4);
        assert_eq!(model.equations.len(), 4);

        // Each species should have an equation
        for species_name in &["A", "B", "C", "D"] {
            let found = model.equations.iter().any(|eq| match &eq.lhs {
                Expr::Operator(node) if node.op == "D" => match &node.args[0] {
                    Expr::Variable(name) => name == species_name,
                    _ => false,
                },
                _ => false,
            });
            assert!(found, "Should have equation for species {species_name}");
        }
    }

    #[test]
    #[cfg(feature = "parallel")]
    fn test_stoichiometric_matrix_parallel() {
        let system = ReactionSystem {
            species: [
                create_test_species("A"),
                create_test_species("B"),
                create_test_species("C"),
            ]
            .into_iter()
            .collect::<IndexMap<_, _>>(),
            parameters: IndexMap::new(),
            reactions: vec![
                // A -> B
                create_test_reaction(
                    vec![("A", 1.0)],
                    vec![("B", 1.0)],
                    Expr::Variable("k1".to_string()),
                ),
                // B -> C
                create_test_reaction(
                    vec![("B", 1.0)],
                    vec![("C", 1.0)],
                    Expr::Variable("k2".to_string()),
                ),
            ],
            ..Default::default()
        };

        // Test the parallel version
        let parallel_matrix = crate::reactions::stoichiometric_matrix_parallel(&system)
            .expect("Parallel stoichiometric matrix should work");

        // Test the sequential version for comparison
        let sequential_matrix = stoichiometric_matrix(&system);

        // Both should produce the same result
        assert_eq!(parallel_matrix.len(), sequential_matrix.len());
        assert_eq!(parallel_matrix[0].len(), sequential_matrix[0].len());

        for (i, row) in parallel_matrix.iter().enumerate() {
            for (j, &val) in row.iter().enumerate() {
                assert_eq!(
                    val, sequential_matrix[i][j],
                    "Parallel and sequential results should match at position [{i}, {j}]"
                );
            }
        }

        // Verify specific values
        assert_eq!(parallel_matrix[0][0], -1.0, "A consumed in reaction 1");
        assert_eq!(parallel_matrix[1][0], 1.0, "B produced in reaction 1");
        assert_eq!(parallel_matrix[1][1], -1.0, "B consumed in reaction 2");
        assert_eq!(parallel_matrix[2][1], 1.0, "C produced in reaction 2");
    }
}