Entropyk/crates/solver/src/system.rs

//! System graph structure for thermodynamic simulation.
//!
//! This module provides the core graph representation of a thermodynamic system,
//! where nodes are components and edges represent flow connections. Edges index
//! into the solver's state vector (P and h per edge).
//!
//! Multi-circuit support (Story 3.3): A machine can have up to 5 independent
//! circuits (valid circuit IDs: 0, 1, 2, 3, 4). Each node belongs to exactly one
//! circuit. Flow edges connect only nodes within the same circuit.

use entropyk_components::{
    validate_port_continuity, Component, ComponentError, ConnectionError, JacobianBuilder,
    ResidualVector, SystemState as StateSlice,
};
use petgraph::algo;
use petgraph::graph::{EdgeIndex, Graph, NodeIndex};
use petgraph::visit::EdgeRef;
use petgraph::Directed;
use std::collections::HashMap;

use crate::coupling::{has_circular_dependencies, ThermalCoupling};
use crate::error::{AddEdgeError, TopologyError};
use entropyk_core::Temperature;

/// Circuit identifier. Valid range 0..=4 (max 5 circuits per machine).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CircuitId(pub u8);

impl CircuitId {
    /// Maximum circuit ID (inclusive). Machine supports up to 5 circuits.
    pub const MAX: u8 = 4;

    /// Creates a new CircuitId if within valid range.
    ///
    /// # Errors
    ///
    /// Returns `TopologyError::TooManyCircuits` if `id > 4`.
    pub fn new(id: u8) -> Result<Self, TopologyError> {
        if id <= Self::MAX {
            Ok(CircuitId(id))
        } else {
            Err(TopologyError::TooManyCircuits { requested: id })
        }
    }

    /// Circuit 0 (default for single-circuit systems).
    pub const ZERO: CircuitId = CircuitId(0);
}

/// Weight for flow edges in the system graph.
///
/// Each edge represents a flow connection between two component ports and stores
/// the state vector indices for pressure (P) and enthalpy (h) at that connection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FlowEdge {
    /// State vector index for pressure (Pa)
    pub state_index_p: usize,
    /// State vector index for enthalpy (J/kg)
    pub state_index_h: usize,
}

/// System graph structure.
///
/// Nodes are components (`Box<dyn Component>`), edges are flow connections with
/// state indices. The state vector layout is:
///
/// ```text
/// [P_edge0, h_edge0, P_edge1, h_edge1, ...]
/// ```
///
/// Edge order follows the graph's internal edge iteration order (stable after
/// `finalize()` is called).
pub struct System {
    graph: Graph<Box<dyn Component>, FlowEdge, Directed>,
    /// Maps EdgeIndex to (state_index_p, state_index_h) - built in finalize()
    edge_to_state: HashMap<EdgeIndex, (usize, usize)>,
    /// Maps NodeIndex to CircuitId. Nodes without entry default to circuit 0.
    node_to_circuit: HashMap<NodeIndex, CircuitId>,
    /// Thermal couplings between circuits (heat transfer without fluid mixing).
    thermal_couplings: Vec<ThermalCoupling>,
    finalized: bool,
}

impl System {
    /// Creates a new empty system graph.
    pub fn new() -> Self {
        Self {
            graph: Graph::new(),
            edge_to_state: HashMap::new(),
            node_to_circuit: HashMap::new(),
            thermal_couplings: Vec::new(),
            finalized: false,
        }
    }

    /// Adds a component as a node in the default circuit (circuit 0) and returns its node index.
    ///
    /// For multi-circuit machines, use [`add_component_to_circuit`](Self::add_component_to_circuit).
    pub fn add_component(&mut self, component: Box<dyn Component>) -> NodeIndex {
        self.add_component_to_circuit(component, CircuitId::ZERO)
            .unwrap()
    }

    /// Adds a component as a node in the specified circuit and returns its node index.
    ///
    /// # Errors
    ///
    /// Returns `TopologyError::TooManyCircuits` if `circuit_id.0 > 4`.
    pub fn add_component_to_circuit(
        &mut self,
        component: Box<dyn Component>,
        circuit_id: CircuitId,
    ) -> Result<NodeIndex, TopologyError> {
        CircuitId::new(circuit_id.0)?;
        self.finalized = false;
        let node_idx = self.graph.add_node(component);
        self.node_to_circuit.insert(node_idx, circuit_id);
        Ok(node_idx)
    }

    /// Returns the circuit ID for a node, or circuit 0 if not found (backward compat).
    pub fn node_circuit(&self, node: NodeIndex) -> CircuitId {
        self.node_to_circuit
            .get(&node)
            .copied()
            .unwrap_or(CircuitId::ZERO)
    }

    /// Returns the circuit ID for an edge based on its source node.
    ///
    /// # Panics
    ///
    /// Panics if the edge index is invalid.
    pub fn edge_circuit(&self, edge: EdgeIndex) -> CircuitId {
        let (src, _tgt) = self.graph.edge_endpoints(edge).expect("invalid edge index");
        self.node_circuit(src)
    }

    /// Adds a flow edge from `source` to `target` without port validation.
    ///
    /// **No port compatibility validation is performed.** Use
    /// [`add_edge_with_ports`](Self::add_edge_with_ports) for components with ports to validate
    /// fluid, pressure, and enthalpy continuity. This method is intended for components without
    /// ports (e.g. mock components in tests).
    ///
    /// Flow edges connect only nodes within the same circuit. Cross-circuit connections
    /// are rejected (thermal coupling is Story 3.4).
    ///
    /// State indices are assigned when `finalize()` is called.
    ///
    /// # Errors
    ///
    /// Returns `TopologyError::CrossCircuitConnection` if source and target are in different circuits.
    pub fn add_edge(
        &mut self,
        source: NodeIndex,
        target: NodeIndex,
    ) -> Result<EdgeIndex, TopologyError> {
        let src_circuit = self.node_circuit(source);
        let tgt_circuit = self.node_circuit(target);
        if src_circuit != tgt_circuit {
            tracing::warn!(
                "Cross-circuit edge rejected: source circuit {}, target circuit {}",
                src_circuit.0,
                tgt_circuit.0
            );
            return Err(TopologyError::CrossCircuitConnection {
                source_circuit: src_circuit.0,
                target_circuit: tgt_circuit.0,
            });
        }

        // Safety check: Warn if connecting components with ports using the non-validating method
        if let (Some(src), Some(tgt)) = (
            self.graph.node_weight(source),
            self.graph.node_weight(target),
        ) {
            if !src.get_ports().is_empty() || !tgt.get_ports().is_empty() {
                tracing::warn!(
                    "add_edge called on components with ports (src: {:?}, tgt: {:?}). \
                     This bypasses port validation. Use add_edge_with_ports instead.",
                    source,
                    target
                );
            }
        }

        self.finalized = false;
        Ok(self.graph.add_edge(
            source,
            target,
            FlowEdge {
                state_index_p: 0,
                state_index_h: 0,
            },
        ))
    }

    /// Adds a flow edge from `source` outlet port to `target` inlet port with validation.
    ///
    /// Validates circuit membership (same circuit), then fluid compatibility, pressure and
    /// enthalpy continuity using port.rs tolerances. For 2-port components: `source_port_idx=1`
    /// (outlet), `target_port_idx=0` (inlet).
    ///
    /// # Errors
    ///
    /// Returns `AddEdgeError::Topology` if source and target are in different circuits.
    /// Returns `AddEdgeError::Connection` if ports are incompatible (fluid, pressure, or enthalpy mismatch).
    pub fn add_edge_with_ports(
        &mut self,
        source: NodeIndex,
        source_port_idx: usize,
        target: NodeIndex,
        target_port_idx: usize,
    ) -> Result<EdgeIndex, AddEdgeError> {
        // Circuit validation first
        let src_circuit = self.node_circuit(source);
        let tgt_circuit = self.node_circuit(target);
        if src_circuit != tgt_circuit {
            tracing::warn!(
                "Cross-circuit edge rejected: source circuit {}, target circuit {}",
                src_circuit.0,
                tgt_circuit.0
            );
            return Err(TopologyError::CrossCircuitConnection {
                source_circuit: src_circuit.0,
                target_circuit: tgt_circuit.0,
            }
            .into());
        }

        let source_comp = self
            .graph
            .node_weight(source)
            .ok_or_else(|| ConnectionError::InvalidNodeIndex(source.index()))?;
        let target_comp = self
            .graph
            .node_weight(target)
            .ok_or_else(|| ConnectionError::InvalidNodeIndex(target.index()))?;

        let source_ports = source_comp.get_ports();
        let target_ports = target_comp.get_ports();

        if source_ports.is_empty() && target_ports.is_empty() {
            // No ports: add edge without validation (backward compat)
            self.finalized = false;
            return Ok(self.graph.add_edge(
                source,
                target,
                FlowEdge {
                    state_index_p: 0,
                    state_index_h: 0,
                },
            ));
        }

        if source_port_idx >= source_ports.len() {
            return Err(ConnectionError::InvalidPortIndex {
                index: source_port_idx,
                port_count: source_ports.len(),
                max_index: source_ports.len().saturating_sub(1),
            }
            .into());
        }
        if target_port_idx >= target_ports.len() {
            return Err(ConnectionError::InvalidPortIndex {
                index: target_port_idx,
                port_count: target_ports.len(),
                max_index: target_ports.len().saturating_sub(1),
            }
            .into());
        }

        let outlet = &source_ports[source_port_idx];
        let inlet = &target_ports[target_port_idx];
        if let Err(ref e) = validate_port_continuity(outlet, inlet) {
            tracing::warn!("Port validation failed: {}", e);
            return Err(e.clone().into());
        }

        self.finalized = false;
        Ok(self.graph.add_edge(
            source,
            target,
            FlowEdge {
                state_index_p: 0,
                state_index_h: 0,
            },
        ))
    }

    /// Finalizes the graph: builds edge→state index mapping and validates topology.
    ///
    /// # State vector layout
    ///
    /// The state vector has length `2 * edge_count`. For each edge (in graph iteration order):
    /// - `state[2*i]` = pressure at edge i (Pa)
    /// - `state[2*i + 1]` = enthalpy at edge i (J/kg)
    ///
    /// # Errors
    ///
    /// Returns `TopologyError` if:
    /// - Any node is isolated (no edges)
    /// - The graph is empty (no components)
    pub fn finalize(&mut self) -> Result<(), TopologyError> {
        self.validate_topology()?;

        if !self.thermal_couplings.is_empty() && has_circular_dependencies(self.thermal_couplings())
        {
            tracing::warn!("Circular thermal coupling detected, simultaneous solving required");
        }

        self.edge_to_state.clear();
        let mut idx = 0usize;
        for edge_idx in self.graph.edge_indices() {
            let (p_idx, h_idx) = (idx, idx + 1);
            self.edge_to_state.insert(edge_idx, (p_idx, h_idx));
            if let Some(weight) = self.graph.edge_weight_mut(edge_idx) {
                weight.state_index_p = p_idx;
                weight.state_index_h = h_idx;
            }
            idx += 2;
        }
        self.finalized = true;
        Ok(())
    }

    /// Validates the topology: no isolated nodes and edge circuit consistency.
    ///
    /// Note: "All ports connected" validation requires port→edge association
    /// (Story 3.2 Port Compatibility Validation).
    fn validate_topology(&self) -> Result<(), TopologyError> {
        let node_count = self.graph.node_count();
        if node_count == 0 {
            return Ok(());
        }

        for node_idx in self.graph.node_indices() {
            let degree = self
                .graph
                .edges_directed(node_idx, petgraph::Direction::Incoming)
                .count()
                + self
                    .graph
                    .edges_directed(node_idx, petgraph::Direction::Outgoing)
                    .count();
            if degree == 0 {
                return Err(TopologyError::IsolatedNode {
                    node_index: node_idx.index(),
                });
            }
        }

        // Validate that all edges connect nodes within the same circuit
        for edge_idx in self.graph.edge_indices() {
            if let Some((src, tgt)) = self.graph.edge_endpoints(edge_idx) {
                let src_circuit = self.node_circuit(src);
                let tgt_circuit = self.node_circuit(tgt);
                if src_circuit != tgt_circuit {
                    return Err(TopologyError::CrossCircuitConnection {
                        source_circuit: src_circuit.0,
                        target_circuit: tgt_circuit.0,
                    });
                }
            }
        }

        Ok(())
    }

    /// Returns the documented state vector layout.
    ///
    /// Layout: `[P_edge0, h_edge0, P_edge1, h_edge1, ...]` where each edge (in
    /// graph iteration order) contributes 2 entries: pressure (Pa) then enthalpy (J/kg).
    ///
    /// # Panics
    ///
    /// Panics if `finalize()` has not been called.
    pub fn state_layout(&self) -> &'static str {
        assert!(self.finalized, "call finalize() before state_layout()");
        "[P_edge0, h_edge0, P_edge1, h_edge1, ...] — 2 per edge (pressure Pa, enthalpy J/kg)"
    }

    /// Returns the length of the state vector: `2 * edge_count`.
    ///
    /// # Panics
    ///
    /// Panics if `finalize()` has not been called.
    pub fn state_vector_len(&self) -> usize {
        assert!(self.finalized, "call finalize() before state_vector_len()");
        2 * self.graph.edge_count()
    }

    /// Returns the state indices (P, h) for the given edge.
    ///
    /// # Returns
    ///
    /// `(state_index_p, state_index_h)` for the edge.
    ///
    /// # Panics
    ///
    /// Panics if `finalize()` has not been called or if `edge_id` is invalid.
    pub fn edge_state_indices(&self, edge_id: EdgeIndex) -> (usize, usize) {
        assert!(
            self.finalized,
            "call finalize() before edge_state_indices()"
        );
        *self
            .edge_to_state
            .get(&edge_id)
            .expect("invalid edge index")
    }

    /// Returns the number of edges in the graph.
    pub fn edge_count(&self) -> usize {
        self.graph.edge_count()
    }

    /// Returns an iterator over all edge indices in the graph.
    pub fn edge_indices(&self) -> impl Iterator<Item = EdgeIndex> + '_ {
        self.graph.edge_indices()
    }

    /// Returns the number of nodes (components) in the graph.
    pub fn node_count(&self) -> usize {
        self.graph.node_count()
    }

    /// Returns the number of distinct circuits in the machine.
    ///
    /// Circuits are identified by the unique circuit IDs present in `node_to_circuit`.
    /// Empty system returns 0. Systems with components always return >= 1 since
    /// all components are assigned to a circuit (defaulting to circuit 0).
    /// Valid circuit IDs are 0 through 4 (inclusive), supporting up to 5 circuits.
    pub fn circuit_count(&self) -> usize {
        if self.graph.node_count() == 0 {
            return 0;
        }
        let mut ids: Vec<u8> = self.node_to_circuit.values().map(|c| c.0).collect();
        if ids.is_empty() {
            // This shouldn't happen since add_component adds to node_to_circuit,
            // but handle defensively
            return 1;
        }
        ids.sort_unstable();
        ids.dedup();
        ids.len()
    }

    /// Returns an iterator over node indices belonging to the given circuit.
    pub fn circuit_nodes(&self, circuit_id: CircuitId) -> impl Iterator<Item = NodeIndex> + '_ {
        self.graph.node_indices().filter(move |&idx| {
            self.node_to_circuit
                .get(&idx)
                .copied()
                .unwrap_or(CircuitId::ZERO)
                == circuit_id
        })
    }

    /// Returns an iterator over edge indices belonging to the given circuit.
    ///
    /// An edge belongs to a circuit if both its source and target nodes are in that circuit.
    pub fn circuit_edges(&self, circuit_id: CircuitId) -> impl Iterator<Item = EdgeIndex> + '_ {
        self.graph.edge_indices().filter(move |&edge_idx| {
            let (src, tgt) = self.graph.edge_endpoints(edge_idx).expect("valid edge");
            self.node_circuit(src) == circuit_id && self.node_circuit(tgt) == circuit_id
        })
    }

    /// Checks if a circuit has any components.
    fn circuit_exists(&self, circuit_id: CircuitId) -> bool {
        self.node_to_circuit.values().any(|&c| c == circuit_id)
    }

    /// Adds a thermal coupling between two circuits.
    ///
    /// Thermal couplings represent heat exchangers that transfer heat from a "hot"
    /// circuit to a "cold" circuit without fluid mixing. Heat flows from hot to cold
    /// proportional to the temperature difference and thermal conductance (UA).
    ///
    /// # Arguments
    ///
    /// * `coupling` - The thermal coupling to add
    ///
    /// # Returns
    ///
    /// The index of the added coupling in the internal storage.
    ///
    /// # Errors
    ///
    /// Returns `TopologyError::InvalidCircuitForCoupling` if either circuit
    /// referenced in the coupling does not exist in the system.
    ///
    /// # Example
    ///
    /// ```no_run
    /// use entropyk_solver::{System, ThermalCoupling, CircuitId};
    /// use entropyk_core::ThermalConductance;
    /// use entropyk_components::Component;
    /// # fn make_mock() -> Box<dyn Component> { unimplemented!() }
    ///
    /// let mut sys = System::new();
    /// sys.add_component_to_circuit(make_mock(), CircuitId(0)).unwrap();
    /// sys.add_component_to_circuit(make_mock(), CircuitId(1)).unwrap();
    ///
    /// let coupling = ThermalCoupling::new(
    ///     CircuitId(0),
    ///     CircuitId(1),
    ///     ThermalConductance::from_watts_per_kelvin(1000.0),
    /// );
    /// let idx = sys.add_thermal_coupling(coupling).unwrap();
    /// ```
    pub fn add_thermal_coupling(
        &mut self,
        coupling: ThermalCoupling,
    ) -> Result<usize, TopologyError> {
        // Validate that both circuits exist
        if !self.circuit_exists(coupling.hot_circuit) {
            return Err(TopologyError::InvalidCircuitForCoupling {
                circuit_id: coupling.hot_circuit.0,
            });
        }
        if !self.circuit_exists(coupling.cold_circuit) {
            return Err(TopologyError::InvalidCircuitForCoupling {
                circuit_id: coupling.cold_circuit.0,
            });
        }

        self.finalized = false;
        self.thermal_couplings.push(coupling);
        Ok(self.thermal_couplings.len() - 1)
    }

    /// Returns the number of thermal couplings in the system.
    pub fn thermal_coupling_count(&self) -> usize {
        self.thermal_couplings.len()
    }

    /// Returns a reference to all thermal couplings.
    pub fn thermal_couplings(&self) -> &[ThermalCoupling] {
        &self.thermal_couplings
    }

    /// Returns a reference to a specific thermal coupling by index.
    pub fn get_thermal_coupling(&self, index: usize) -> Option<&ThermalCoupling> {
        self.thermal_couplings.get(index)
    }

    /// Returns the number of coupling residual equations (one per thermal coupling).
    ///
    /// The solver must reserve this many rows in the residual vector for coupling
    /// heat balance equations. Use [`coupling_residuals`](Self::coupling_residuals) to fill them.
    pub fn coupling_residual_count(&self) -> usize {
        self.thermal_couplings.len()
    }

    /// Fills coupling residuals into `out`.
    ///
    /// For each thermal coupling, the residual is the heat transfer rate Q (W) into the cold
    /// circuit: Q = η·UA·(T_hot − T_cold). The solver typically uses this in a heat balance
    /// (e.g. r = Q_actual − Q_expected). Temperatures must be in Kelvin.
    ///
    /// # Arguments
    ///
    /// * `temperatures` - One (T_hot_K, T_cold_K) per coupling; length must equal
    ///   `thermal_coupling_count()`. The solver obtains these from state (e.g. P, h → T via fluid backend).
    /// * `out` - Slice to write residuals; length must be at least `coupling_residual_count()`.
    pub fn coupling_residuals(&self, temperatures: &[(f64, f64)], out: &mut [f64]) {
        assert!(
            temperatures.len() == self.thermal_couplings.len(),
            "temperatures.len() must equal thermal_coupling_count()"
        );
        assert!(
            out.len() >= self.thermal_couplings.len(),
            "out.len() must be at least coupling_residual_count()"
        );
        for (i, coupling) in self.thermal_couplings.iter().enumerate() {
            let (t_hot_k, t_cold_k) = temperatures[i];
            let t_hot = Temperature::from_kelvin(t_hot_k);
            let t_cold = Temperature::from_kelvin(t_cold_k);
            out[i] = crate::coupling::compute_coupling_heat(coupling, t_hot, t_cold);
        }
    }

    /// Returns Jacobian entries for coupling residuals with respect to temperature state.
    ///
    /// The solver state may include temperature unknowns for coupling interfaces (or T derived from P, h).
    /// For each coupling i: ∂Q_i/∂T_hot = η·UA, ∂Q_i/∂T_cold = −η·UA.
    ///
    /// # Arguments
    ///
    /// * `row_offset` - First row index for coupling equations in the global residual vector.
    /// * `t_hot_cols` - State column index for T_hot per coupling; length = `thermal_coupling_count()`.
    /// * `t_cold_cols` - State column index for T_cold per coupling; length = `thermal_coupling_count()`.
    ///
    /// # Returns
    ///
    /// `(row, col, value)` tuples for the Jacobian. Row is `row_offset + coupling_index`.
    pub fn coupling_jacobian_entries(
        &self,
        row_offset: usize,
        t_hot_cols: &[usize],
        t_cold_cols: &[usize],
    ) -> Vec<(usize, usize, f64)> {
        assert!(
            t_hot_cols.len() == self.thermal_couplings.len()
                && t_cold_cols.len() == self.thermal_couplings.len(),
            "t_hot_cols and t_cold_cols length must equal thermal_coupling_count()"
        );
        let mut entries = Vec::with_capacity(2 * self.thermal_couplings.len());
        for (i, coupling) in self.thermal_couplings.iter().enumerate() {
            let dr_dt_hot = coupling.efficiency * coupling.ua.to_watts_per_kelvin();
            let dr_dt_cold = -dr_dt_hot;
            let row = row_offset + i;
            entries.push((row, t_hot_cols[i], dr_dt_hot));
            entries.push((row, t_cold_cols[i], dr_dt_cold));
        }
        entries
    }

    /// Returns true if the graph contains a cycle.
    ///
    /// Refrigeration circuits form cycles (compressor → condenser → valve → evaporator → compressor),
    /// so cycles are expected and valid.
    pub fn is_cyclic(&self) -> bool {
        algo::is_cyclic_directed(&self.graph)
    }

    /// Iterates over (component, edge_indices) for Jacobian assembly.
    ///
    /// For each node, yields the component and a map from edge index to (state_index_p, state_index_h)
    /// for edges incident to that node (incoming and outgoing).
    ///
    /// # Panics
    ///
    /// Panics if `finalize()` has not been called.
    pub fn traverse_for_jacobian(
        &self,
    ) -> impl Iterator<Item = (NodeIndex, &dyn Component, Vec<(EdgeIndex, usize, usize)>)> {
        assert!(
            self.finalized,
            "call finalize() before traverse_for_jacobian()"
        );

        self.graph.node_indices().map(move |node_idx| {
            let component = self.graph.node_weight(node_idx).unwrap();
            let mut edge_indices = Vec::new();

            for edge_ref in self
                .graph
                .edges_directed(node_idx, petgraph::Direction::Incoming)
            {
                let edge_idx = edge_ref.id();
                if let Some(&(p, h)) = self.edge_to_state.get(&edge_idx) {
                    edge_indices.push((edge_idx, p, h));
                }
            }
            for edge_ref in self
                .graph
                .edges_directed(node_idx, petgraph::Direction::Outgoing)
            {
                let edge_idx = edge_ref.id();
                if let Some(&(p, h)) = self.edge_to_state.get(&edge_idx) {
                    edge_indices.push((edge_idx, p, h));
                }
            }

            (node_idx, component.as_ref(), edge_indices)
        })
    }

    /// Assembles residuals from all components.
    ///
    /// Components receive the full state slice and write to their equation indices.
    /// Equation indices are computed from component order and `n_equations()`.
    ///
    /// # Errors
    ///
    /// Returns `ComponentError::InvalidResidualDimensions` if `residuals.len()` is
    /// less than the total number of equations across all components.
    pub fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        let total_eqs: usize = self
            .traverse_for_jacobian()
            .map(|(_, c, _)| c.n_equations())
            .sum();
        if residuals.len() < total_eqs {
            return Err(ComponentError::InvalidResidualDimensions {
                expected: total_eqs,
                actual: residuals.len(),
            });
        }

        let mut eq_offset = 0;
        for (_node_idx, component, _edge_indices) in self.traverse_for_jacobian() {
            let n = component.n_equations();
            if n > 0 {
                let mut temp = vec![0.0; n];
                component.compute_residuals(state, &mut temp)?;
                residuals[eq_offset..eq_offset + n].copy_from_slice(&temp);
            }
            eq_offset += n;
        }
        Ok(())
    }

    /// Assembles Jacobian entries from all components.
    ///
    /// Each component receives the state and writes to JacobianBuilder. The
    /// [`traverse_for_jacobian`](Self::traverse_for_jacobian) iterator provides
    /// the edge→state mapping `(EdgeIndex, state_index_p, state_index_h)` per
    /// component. Components must know their port→state mapping (e.g. from graph
    /// construction in Story 3.2) to write correct column indices.
    pub fn assemble_jacobian(
        &self,
        state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        let mut row_offset = 0;
        for (_node_idx, component, _edge_indices) in self.traverse_for_jacobian() {
            let n = component.n_equations();
            if n > 0 {
                // Components write rows 0..n-1; we offset to global equation indices.
                let mut temp_builder = JacobianBuilder::new();
                component.jacobian_entries(state, &mut temp_builder)?;
                for (r, c, v) in temp_builder.entries() {
                    jacobian.add_entry(row_offset + r, *c, *v);
                }
            }
            row_offset += n;
        }
        Ok(())
    }
}

impl Default for System {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;
    use entropyk_components::port::{FluidId, Port};
    use entropyk_components::{ConnectedPort, SystemState};
    use entropyk_core::{Enthalpy, Pressure};

    /// Minimal mock component for testing.
    struct MockComponent {
        n_equations: usize,
    }

    impl Component for MockComponent {
        fn compute_residuals(
            &self,
            _state: &SystemState,
            residuals: &mut entropyk_components::ResidualVector,
        ) -> Result<(), ComponentError> {
            for r in residuals.iter_mut().take(self.n_equations) {
                *r = 0.0;
            }
            Ok(())
        }

        fn jacobian_entries(
            &self,
            _state: &SystemState,
            jacobian: &mut JacobianBuilder,
        ) -> Result<(), ComponentError> {
            for i in 0..self.n_equations {
                jacobian.add_entry(i, i, 1.0);
            }
            Ok(())
        }

        fn n_equations(&self) -> usize {
            self.n_equations
        }

        fn get_ports(&self) -> &[ConnectedPort] {
            &[]
        }
    }

    fn make_mock(n: usize) -> Box<dyn Component> {
        Box::new(MockComponent { n_equations: n })
    }

    /// Mock component with 2 ports (inlet=0, outlet=1) for port validation tests.
    fn make_ported_mock(fluid: &str, pressure_pa: f64, enthalpy_jkg: f64) -> Box<dyn Component> {
        let inlet = Port::new(
            FluidId::new(fluid),
            Pressure::from_pascals(pressure_pa),
            Enthalpy::from_joules_per_kg(enthalpy_jkg),
        );
        let outlet = Port::new(
            FluidId::new(fluid),
            Pressure::from_pascals(pressure_pa),
            Enthalpy::from_joules_per_kg(enthalpy_jkg),
        );
        let (connected_inlet, connected_outlet) = inlet.connect(outlet).unwrap();
        let ports: Vec<ConnectedPort> = vec![connected_inlet, connected_outlet];
        Box::new(PortedMockComponent { ports })
    }

    struct PortedMockComponent {
        ports: Vec<ConnectedPort>,
    }

    impl Component for PortedMockComponent {
        fn compute_residuals(
            &self,
            _state: &SystemState,
            residuals: &mut entropyk_components::ResidualVector,
        ) -> Result<(), ComponentError> {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            Ok(())
        }

        fn jacobian_entries(
            &self,
            _state: &SystemState,
            _jacobian: &mut JacobianBuilder,
        ) -> Result<(), ComponentError> {
            Ok(())
        }

        fn n_equations(&self) -> usize {
            0
        }

        fn get_ports(&self) -> &[ConnectedPort] {
            &self.ports
        }
    }

    #[test]
    fn test_simple_cycle_builds() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(0));
        let n1 = sys.add_component(make_mock(0));
        let n2 = sys.add_component(make_mock(0));
        let n3 = sys.add_component(make_mock(0));

        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n2).unwrap();
        sys.add_edge(n2, n3).unwrap();
        sys.add_edge(n3, n0).unwrap();

        assert_eq!(sys.node_count(), 4);
        assert_eq!(sys.edge_count(), 4);

        let result = sys.finalize();
        assert!(
            result.is_ok(),
            "finalize should succeed: {:?}",
            result.err()
        );
    }

    #[test]
    fn test_state_vector_length() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(0));
        let n1 = sys.add_component(make_mock(0));
        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n0).unwrap();

        sys.finalize().unwrap();
        assert_eq!(sys.state_vector_len(), 4); // 2 edges * 2 = 4
    }

    #[test]
    fn test_edge_indices_contiguous() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(0));
        let n1 = sys.add_component(make_mock(0));
        let n2 = sys.add_component(make_mock(0));

        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n2).unwrap();
        sys.add_edge(n2, n0).unwrap();

        sys.finalize().unwrap();

        let mut indices: Vec<usize> = Vec::new();
        for edge_idx in sys.edge_indices() {
            let (p, h) = sys.edge_state_indices(edge_idx);
            indices.push(p);
            indices.push(h);
        }

        let n = sys.edge_count();
        assert_eq!(indices.len(), 2 * n);
        let expected: Vec<usize> = (0..2 * n).collect();
        assert_eq!(indices, expected, "indices should be 0..2n");
    }

    #[test]
    fn test_cycle_detected() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(0));
        let n1 = sys.add_component(make_mock(0));
        let n2 = sys.add_component(make_mock(0));

        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n2).unwrap();
        sys.add_edge(n2, n0).unwrap();

        sys.finalize().unwrap();
        assert!(
            sys.is_cyclic(),
            "refrigeration cycle should be detected as cyclic"
        );
    }

    #[test]
    fn test_dangling_node_error() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(0));
        let n1 = sys.add_component(make_mock(0));
        let n2 = sys.add_component(make_mock(0)); // isolated

        sys.add_edge(n0, n1).unwrap();
        // n2 has no edges

        let result = sys.finalize();
        assert!(result.is_err());
        match &result {
            Err(TopologyError::IsolatedNode { node_index }) => {
                assert!(
                    *node_index < sys.node_count(),
                    "isolated node index {} must be < node_count {}",
                    node_index,
                    sys.node_count()
                );
                assert_eq!(*node_index, n2.index(), "isolated node should be n2");
            }
            other => panic!("expected IsolatedNode error, got {:?}", other),
        }
    }

    #[test]
    fn test_traverse_components() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(1));
        let n1 = sys.add_component(make_mock(1));
        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n0).unwrap();

        sys.finalize().unwrap();

        let mut count = 0;
        for (_node_idx, component, edge_indices) in sys.traverse_for_jacobian() {
            count += 1;
            assert_eq!(component.n_equations(), 1);
            assert_eq!(
                edge_indices.len(),
                2,
                "each node has 2 incident edges in 2-node cycle"
            );
            for (_edge_idx, p, h) in &edge_indices {
                assert!(p < &sys.state_vector_len());
                assert!(h < &sys.state_vector_len());
                assert_eq!(h, &(p + 1));
            }
        }
        assert_eq!(count, 2);
    }

    #[test]
    fn test_empty_graph_finalize_ok() {
        let mut sys = System::new();
        let result = sys.finalize();
        assert!(result.is_ok());
    }

    #[test]
    fn test_state_layout_integration() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(1));
        let n1 = sys.add_component(make_mock(1));
        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n0).unwrap();

        sys.finalize().unwrap();

        let layout = sys.state_layout();
        assert!(layout.contains("P_edge"));
        assert!(layout.contains("h_edge"));

        let state_len = sys.state_vector_len();
        assert_eq!(state_len, 4);

        let mut state = vec![0.0; state_len];
        state[0] = 1e5; // P_edge0
        state[1] = 250000.0; // h_edge0
        state[2] = 5e5; // P_edge1
        state[3] = 300000.0; // h_edge1

        let mut residuals = vec![0.0; 2];
        let result = sys.compute_residuals(&state, &mut residuals);
        assert!(result.is_ok());
    }

    #[test]
    fn test_valid_connection_same_fluid() {
        let p = 100_000.0;
        let h = 400_000.0;
        let mut sys = System::new();
        let n0 = sys.add_component(make_ported_mock("R134a", p, h));
        let n1 = sys.add_component(make_ported_mock("R134a", p, h));

        let result = sys.add_edge_with_ports(n0, 1, n1, 0);
        assert!(
            result.is_ok(),
            "R134a to R134a should succeed: {:?}",
            result.err()
        );
        sys.add_edge(n1, n0).unwrap(); // backward edge, no ports
        sys.finalize().unwrap();
    }

    #[test]
    fn test_incompatible_fluid_rejected() {
        let p = 100_000.0;
        let h = 400_000.0;
        let mut sys = System::new();
        let n0 = sys.add_component(make_ported_mock("R134a", p, h));
        let n1 = sys.add_component(make_ported_mock("Water", p, h));

        let result = sys.add_edge_with_ports(n0, 1, n1, 0);
        assert!(result.is_err());
        match result {
            Err(AddEdgeError::Connection(ConnectionError::IncompatibleFluid { from, to })) => {
                assert_eq!(from, "R134a");
                assert_eq!(to, "Water");
            }
            other => panic!(
                "expected AddEdgeError::Connection(IncompatibleFluid), got {:?}",
                other
            ),
        }
    }

    #[test]
    fn test_pressure_mismatch_rejected() {
        let h = 400_000.0;
        let mut sys = System::new();
        let n0 = sys.add_component(make_ported_mock("R134a", 100_000.0, h));
        let n1 = sys.add_component(make_ported_mock("R134a", 200_000.0, h));

        let result = sys.add_edge_with_ports(n0, 1, n1, 0);
        assert!(result.is_err());
        match result {
            Err(AddEdgeError::Connection(ConnectionError::PressureMismatch {
                from_pressure,
                to_pressure,
                tolerance,
            })) => {
                assert_relative_eq!(from_pressure, 100_000.0, epsilon = 1.0);
                assert_relative_eq!(to_pressure, 200_000.0, epsilon = 1.0);
                assert!(tolerance >= 1.0, "tolerance should be at least 1 Pa");
            }
            other => panic!(
                "expected AddEdgeError::Connection(PressureMismatch), got {:?}",
                other
            ),
        }
    }

    #[test]
    fn test_enthalpy_mismatch_rejected() {
        let p = 100_000.0;
        let mut sys = System::new();
        let n0 = sys.add_component(make_ported_mock("R134a", p, 400_000.0));
        let n1 = sys.add_component(make_ported_mock("R134a", p, 500_000.0));

        let result = sys.add_edge_with_ports(n0, 1, n1, 0);
        assert!(result.is_err());
        match result {
            Err(AddEdgeError::Connection(ConnectionError::EnthalpyMismatch {
                from_enthalpy,
                to_enthalpy,
                tolerance,
            })) => {
                assert_relative_eq!(from_enthalpy, 400_000.0, epsilon = 1.0);
                assert_relative_eq!(to_enthalpy, 500_000.0, epsilon = 1.0);
                assert_relative_eq!(tolerance, 100.0, epsilon = 1.0); // ENTHALPY_TOLERANCE_J_KG
            }
            other => panic!(
                "expected AddEdgeError::Connection(EnthalpyMismatch), got {:?}",
                other
            ),
        }
    }

    #[test]
    fn test_pressure_tolerance_boundary() {
        let h = 400_000.0;
        let base_pressure = 100_000.0; // 100 kPa
        let tolerance: f64 = (base_pressure * 1e-4f64).max(1.0f64); // 10 Pa for 100 kPa

        let mut sys = System::new();
        let n0 = sys.add_component(make_ported_mock("R134a", base_pressure, h));

        // Exactly at tolerance - should succeed
        let n1 = sys.add_component(make_ported_mock("R134a", base_pressure + tolerance, h));
        let result = sys.add_edge_with_ports(n0, 1, n1, 0);
        assert!(
            result.is_ok(),
            "Connection at exact tolerance ({:.1} Pa diff) should succeed",
            tolerance
        );

        // Just outside tolerance - should fail
        let n2 = sys.add_component(make_ported_mock(
            "R134a",
            base_pressure + tolerance + 1.0,
            h,
        ));
        let result = sys.add_edge_with_ports(n0, 1, n2, 0);
        assert!(
            result.is_err(),
            "Connection just outside tolerance ({:.1} Pa diff) should fail",
            tolerance + 1.0
        );
        match result {
            Err(AddEdgeError::Connection(ConnectionError::PressureMismatch {
                from_pressure,
                to_pressure,
                tolerance: tol,
            })) => {
                assert_relative_eq!(from_pressure, base_pressure, epsilon = 0.1);
                assert_relative_eq!(to_pressure, base_pressure + tolerance + 1.0, epsilon = 0.1);
                assert_relative_eq!(tol, tolerance, epsilon = 0.1);
            }
            other => panic!("expected PressureMismatch at boundary, got {:?}", other),
        }
    }

    #[test]
    fn test_invalid_port_index_rejected() {
        let p = 100_000.0;
        let h = 400_000.0;
        let mut sys = System::new();
        let n0 = sys.add_component(make_ported_mock("R134a", p, h));
        let n1 = sys.add_component(make_ported_mock("R134a", p, h));

        // Port index 2 out of bounds for 2-port component
        let result = sys.add_edge_with_ports(n0, 2, n1, 0);
        assert!(result.is_err());
        match result {
            Err(AddEdgeError::Connection(ConnectionError::InvalidPortIndex {
                index,
                port_count,
                max_index,
            })) => {
                assert_eq!(index, 2);
                assert_eq!(port_count, 2);
                assert_eq!(max_index, 1);
            }
            other => panic!("expected InvalidPortIndex for source, got {:?}", other),
        }

        // Target port index out of bounds
        let result = sys.add_edge_with_ports(n0, 1, n1, 5);
        assert!(result.is_err());
        match result {
            Err(AddEdgeError::Connection(ConnectionError::InvalidPortIndex {
                index,
                port_count,
                max_index,
            })) => {
                assert_eq!(index, 5);
                assert_eq!(port_count, 2);
                assert_eq!(max_index, 1);
            }
            other => panic!("expected InvalidPortIndex for target, got {:?}", other),
        }
    }

    #[test]
    fn test_simple_cycle_port_validation() {
        let p = 100_000.0;
        let h = 400_000.0;
        let mut sys = System::new();
        let n0 = sys.add_component(make_ported_mock("R134a", p, h));
        let n1 = sys.add_component(make_ported_mock("R134a", p, h));
        let n2 = sys.add_component(make_ported_mock("R134a", p, h));
        let n3 = sys.add_component(make_ported_mock("R134a", p, h));

        sys.add_edge_with_ports(n0, 1, n1, 0).unwrap();
        sys.add_edge_with_ports(n1, 1, n2, 0).unwrap();
        sys.add_edge_with_ports(n2, 1, n3, 0).unwrap();
        sys.add_edge_with_ports(n3, 1, n0, 0).unwrap();

        assert_eq!(sys.edge_count(), 4);
        sys.finalize().unwrap();
    }

    #[test]
    fn test_compute_residuals_bounds_check() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(2));
        let n1 = sys.add_component(make_mock(2));
        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n0).unwrap();
        sys.finalize().unwrap();

        let state = vec![0.0; sys.state_vector_len()];
        let mut residuals = vec![0.0; 1]; // Too small: need 4
        let result = sys.compute_residuals(&state, &mut residuals);
        assert!(result.is_err());
        match result {
            Err(ComponentError::InvalidResidualDimensions { expected, actual }) => {
                assert_eq!(expected, 4);
                assert_eq!(actual, 1);
            }
            other => panic!("expected InvalidResidualDimensions, got {:?}", other),
        }
    }

    // --- Story 3.3: Multi-Circuit Machine Definition tests ---

    #[test]
    fn test_two_circuit_machine() {
        let mut sys = System::new();
        let c0 = CircuitId::ZERO;
        let c1 = CircuitId(1);

        let n0 = sys.add_component_to_circuit(make_mock(0), c0).unwrap();
        let n1 = sys.add_component_to_circuit(make_mock(0), c0).unwrap();
        let n2 = sys.add_component_to_circuit(make_mock(0), c1).unwrap();
        let n3 = sys.add_component_to_circuit(make_mock(0), c1).unwrap();

        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n0).unwrap();
        sys.add_edge(n2, n3).unwrap();
        sys.add_edge(n3, n2).unwrap();

        assert_eq!(sys.circuit_count(), 2);
        assert_eq!(sys.circuit_nodes(c0).count(), 2);
        assert_eq!(sys.circuit_nodes(c1).count(), 2);
        assert_eq!(sys.circuit_edges(c0).count(), 2);
        assert_eq!(sys.circuit_edges(c1).count(), 2);

        sys.finalize().unwrap();
    }

    #[test]
    fn test_cross_circuit_edge_rejected() {
        let mut sys = System::new();
        let c0 = CircuitId::ZERO;
        let c1 = CircuitId(1);

        let n0 = sys.add_component_to_circuit(make_mock(0), c0).unwrap();
        let n1 = sys.add_component_to_circuit(make_mock(0), c1).unwrap();

        let result = sys.add_edge(n0, n1);
        assert!(result.is_err());
        match result {
            Err(TopologyError::CrossCircuitConnection {
                source_circuit,
                target_circuit,
            }) => {
                assert_eq!(source_circuit, 0);
                assert_eq!(target_circuit, 1);
            }
            other => panic!("expected CrossCircuitConnection, got {:?}", other),
        }
    }

    #[test]
    fn test_circuit_count_and_accessors() {
        let mut sys = System::new();
        let c0 = CircuitId::ZERO;
        let c1 = CircuitId(1);
        let c2 = CircuitId(2);

        let _n0 = sys.add_component_to_circuit(make_mock(0), c0).unwrap();
        let _n1 = sys.add_component_to_circuit(make_mock(0), c0).unwrap();
        let _n2 = sys.add_component_to_circuit(make_mock(0), c1).unwrap();
        let _n3 = sys.add_component_to_circuit(make_mock(0), c2).unwrap();

        assert_eq!(sys.circuit_count(), 3);
        assert_eq!(sys.circuit_nodes(c0).count(), 2);
        assert_eq!(sys.circuit_nodes(c1).count(), 1);
        assert_eq!(sys.circuit_nodes(c2).count(), 1);
    }

    #[test]
    fn test_max_five_circuits() {
        // Test: 5 circuits accepted (0, 1, 2, 3, 4), 6th circuit (5) rejected
        let mut sys = System::new();
        for i in 0..=4 {
            let cid = CircuitId(i);
            let result = sys.add_component_to_circuit(make_mock(0), cid);
            assert!(
                result.is_ok(),
                "circuit {} should be accepted (max 5 circuits: 0-4)",
                i
            );
        }
        assert_eq!(sys.circuit_count(), 5, "should have exactly 5 circuits");

        // 6th circuit should be rejected
        let result = sys.add_component_to_circuit(make_mock(0), CircuitId(5));
        assert!(
            result.is_err(),
            "circuit 5 should be rejected (exceeds max of 4)"
        );
        match result {
            Err(TopologyError::TooManyCircuits { requested }) => assert_eq!(requested, 5),
            other => panic!("expected TooManyCircuits, got {:?}", other),
        }
    }

    #[test]
    fn test_single_circuit_backward_compat() {
        let mut sys = System::new();
        let n0 = sys.add_component(make_mock(0));
        let n1 = sys.add_component(make_mock(0));
        let edge0 = sys.add_edge(n0, n1).unwrap();
        let edge1 = sys.add_edge(n1, n0).unwrap();

        assert_eq!(sys.circuit_count(), 1);
        assert_eq!(sys.circuit_nodes(CircuitId::ZERO).count(), 2);

        // Verify edge circuit membership
        assert_eq!(sys.edge_circuit(edge0), CircuitId::ZERO);
        assert_eq!(sys.edge_circuit(edge1), CircuitId::ZERO);

        // Verify circuit_edges returns correct edges
        let circuit_0_edges: Vec<_> = sys.circuit_edges(CircuitId::ZERO).collect();
        assert_eq!(circuit_0_edges.len(), 2);
        assert!(circuit_0_edges.contains(&edge0));
        assert!(circuit_0_edges.contains(&edge1));

        sys.finalize().unwrap();
    }

    #[test]
    fn test_cross_circuit_add_edge_with_ports_rejected() {
        let p = 100_000.0;
        let h = 400_000.0;
        let mut sys = System::new();
        let n0 = sys
            .add_component_to_circuit(make_ported_mock("R134a", p, h), CircuitId::ZERO)
            .unwrap();
        let n1 = sys
            .add_component_to_circuit(make_ported_mock("R134a", p, h), CircuitId(1))
            .unwrap();

        let result = sys.add_edge_with_ports(n0, 1, n1, 0);
        assert!(result.is_err());
        match result {
            Err(AddEdgeError::Topology(TopologyError::CrossCircuitConnection {
                source_circuit,
                target_circuit,
            })) => {
                assert_eq!(source_circuit, 0);
                assert_eq!(target_circuit, 1);
            }
            other => panic!(
                "expected AddEdgeError::Topology(CrossCircuitConnection), got {:?}",
                other
            ),
        }
    }

    // --- Story 3.4: Thermal Coupling Between Circuits tests ---

    #[test]
    fn test_add_thermal_coupling_valid() {
        use entropyk_core::ThermalConductance;

        let mut sys = System::new();
        let _n0 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(0))
            .unwrap();
        let _n1 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(1))
            .unwrap();

        let coupling = ThermalCoupling::new(
            CircuitId(0),
            CircuitId(1),
            ThermalConductance::from_watts_per_kelvin(1000.0),
        );

        let idx = sys.add_thermal_coupling(coupling).unwrap();
        assert_eq!(idx, 0);
        assert_eq!(sys.thermal_coupling_count(), 1);

        let retrieved = sys.get_thermal_coupling(0).unwrap();
        assert_eq!(retrieved.hot_circuit, CircuitId(0));
        assert_eq!(retrieved.cold_circuit, CircuitId(1));
    }

    #[test]
    fn test_add_thermal_coupling_invalid_circuit() {
        use entropyk_core::ThermalConductance;

        let mut sys = System::new();
        let _n0 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(0))
            .unwrap();
        // Circuit 1 has no components

        let coupling = ThermalCoupling::new(
            CircuitId(0),
            CircuitId(1), // This circuit doesn't exist
            ThermalConductance::from_watts_per_kelvin(1000.0),
        );

        let result = sys.add_thermal_coupling(coupling);
        assert!(result.is_err());
        match result {
            Err(TopologyError::InvalidCircuitForCoupling { circuit_id }) => {
                assert_eq!(circuit_id, 1);
            }
            other => panic!("expected InvalidCircuitForCoupling, got {:?}", other),
        }
    }

    #[test]
    fn test_add_thermal_coupling_hot_circuit_invalid() {
        use entropyk_core::ThermalConductance;

        let mut sys = System::new();
        let _n0 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(0))
            .unwrap();

        let coupling = ThermalCoupling::new(
            CircuitId(99), // This circuit doesn't exist
            CircuitId(0),
            ThermalConductance::from_watts_per_kelvin(1000.0),
        );

        let result = sys.add_thermal_coupling(coupling);
        assert!(result.is_err());
        match result {
            Err(TopologyError::InvalidCircuitForCoupling { circuit_id }) => {
                assert_eq!(circuit_id, 99);
            }
            other => panic!("expected InvalidCircuitForCoupling, got {:?}", other),
        }
    }

    #[test]
    fn test_multiple_thermal_couplings() {
        use entropyk_core::ThermalConductance;

        let mut sys = System::new();
        let _n0 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(0))
            .unwrap();
        let _n1 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(1))
            .unwrap();
        let _n2 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(2))
            .unwrap();

        let coupling1 = ThermalCoupling::new(
            CircuitId(0),
            CircuitId(1),
            ThermalConductance::from_watts_per_kelvin(1000.0),
        );
        let coupling2 = ThermalCoupling::new(
            CircuitId(1),
            CircuitId(2),
            ThermalConductance::from_watts_per_kelvin(500.0),
        );

        let idx1 = sys.add_thermal_coupling(coupling1).unwrap();
        let idx2 = sys.add_thermal_coupling(coupling2).unwrap();

        assert_eq!(idx1, 0);
        assert_eq!(idx2, 1);
        assert_eq!(sys.thermal_coupling_count(), 2);

        let all_couplings = sys.thermal_couplings();
        assert_eq!(all_couplings.len(), 2);
    }

    #[test]
    fn test_thermal_coupling_same_circuit() {
        // It's valid to couple a circuit to itself (internal heat exchanger / economizer)
        use entropyk_core::ThermalConductance;

        let mut sys = System::new();
        let _n0 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(0))
            .unwrap();
        let _n1 = sys
            .add_component_to_circuit(make_mock(0), CircuitId(0))
            .unwrap();

        let coupling = ThermalCoupling::new(
            CircuitId(0),
            CircuitId(0), // Same circuit
            ThermalConductance::from_watts_per_kelvin(1000.0),
        );

        let result = sys.add_thermal_coupling(coupling);
        assert!(result.is_ok(), "Same-circuit coupling should be allowed");
    }

    // --- Story 3.5: Zero-Flow Branch Handling ---

    /// Mock component that behaves like an Off branch: residual = state[0] (ṁ - 0), Jacobian ∂r/∂state[0] = 1.
    /// At state[0] = 0 (zero flow) residuals and Jacobian remain finite (no division by zero).
    struct ZeroFlowMock;

    impl Component for ZeroFlowMock {
        fn compute_residuals(
            &self,
            state: &SystemState,
            residuals: &mut entropyk_components::ResidualVector,
        ) -> Result<(), ComponentError> {
            if !state.is_empty() {
                residuals[0] = state[0];
            }
            Ok(())
        }

        fn jacobian_entries(
            &self,
            _state: &SystemState,
            jacobian: &mut JacobianBuilder,
        ) -> Result<(), ComponentError> {
            jacobian.add_entry(0, 0, 1.0);
            Ok(())
        }

        fn n_equations(&self) -> usize {
            1
        }

        fn get_ports(&self) -> &[ConnectedPort] {
            &[]
        }
    }

    #[test]
    fn test_zero_flow_branch_residuals_and_jacobian_finite() {
        // Story 3.5: System with one branch at zero flow must not produce NaN/Inf in residuals or Jacobian.
        let mut sys = System::new();
        let n0 = sys.add_component(Box::new(ZeroFlowMock));
        let n1 = sys.add_component(make_mock(1));
        sys.add_edge(n0, n1).unwrap();
        sys.add_edge(n1, n0).unwrap();
        sys.finalize().unwrap();

        let state_len = sys.state_vector_len();
        let mut state = vec![0.0; state_len];
        state[0] = 0.0;

        let total_eqs: usize = 1 + 1;
        let mut residuals = vec![0.0; total_eqs];
        let result = sys.compute_residuals(&state, &mut residuals);
        assert!(result.is_ok());
        for (i, &r) in residuals.iter().enumerate() {
            assert!(r.is_finite(), "residual[{}] must be finite, got {}", i, r);
        }

        let mut jacobian = JacobianBuilder::new();
        let result = sys.assemble_jacobian(&state, &mut jacobian);
        assert!(result.is_ok());

        let entries = jacobian.entries();
        for (row, col, value) in entries {
            assert!(
                value.is_finite(),
                "Jacobian ({}, {}) must be finite, got {}",
                row,
                col,
                value
            );
        }

        // Check for zero rows: each equation should have at least one non-zero derivative
        let mut row_has_nonzero = vec![false; total_eqs];
        for (row, _col, value) in entries {
            if value.abs() > 1e-15 {
                row_has_nonzero[*row] = true;
            }
        }
        for (row, &has_nonzero) in row_has_nonzero.iter().enumerate() {
            assert!(
                has_nonzero,
                "Jacobian row {} is all zeros (degenerate equation)",
                row
            );
        }
    }
}