//! Integration test: Air-Cooled Chiller with Screw Economizer Compressor
//!
//! Simulates a 2-circuit air-cooled chiller with:
//!   - 2 × ScrewEconomizerCompressor (R134a, VFD controlled 25–60 Hz)
//!   - 4 × MchxCondenserCoil + fan banks (35°C ambient air)
//!   - 2 × FloodedEvaporator + Drum (water-glycol MEG 35%, 12°C → 7°C)
//!   - Economizer (flash-gas injection)
//!   - Superheat control via Constraint
//!   - Fan speed control (anti-override) via BoundedVariable
//!
//! ## Topology per circuit (× 2 circuits)
//!
//! ```text
//!     BrineSource(MEG35%, 12°C)
//!           ↓
//!     FloodedEvaporator ←── Drum ←── Economizer(flash)
//!           ↓                              ↑
//!     ScrewEconomizerCompressor(eco port) ──┘
//!           ↓
//!     FlowSplitter (1 → 2 coils)
//!           ↓                ↓
//!    MchxCoil_A+Fan_A   MchxCoil_B+Fan_B
//!           ↓                ↓
//!     FlowMerger (2 → 1)
//!           ↓
//!     ExpansionValve
//!           ↓
//!     BrineSink(MEG35%, 7°C)
//! ```
//!
//! This test validates topology construction, finalization, and that all
//! components can compute residuals without errors at a reasonable initial state.

use entropyk_components::port::{Connected, FluidId, Port};
use entropyk_components::state_machine::{CircuitId, OperationalState};
use entropyk_components::{
    Component, ComponentError, ConnectedPort, JacobianBuilder, MchxCondenserCoil, Polynomial2D,
    ResidualVector, ScrewEconomizerCompressor, ScrewPerformanceCurves, StateManageable, StateSlice,
};
use entropyk_core::{Enthalpy, MassFlow, Power, Pressure};
use entropyk_solver::{system::System, TopologyError};

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────

#[allow(dead_code)] // Convenience alias kept for readability in this fixture.
type CP = Port<Connected>;

/// Creates a connected port pair — returns the first (connected) port.
fn make_port(fluid: &str, p_bar: f64, h_kj_kg: f64) -> ConnectedPort {
    let a = Port::new(
        FluidId::new(fluid),
        Pressure::from_bar(p_bar),
        Enthalpy::from_joules_per_kg(h_kj_kg * 1000.0),
    );
    let b = Port::new(
        FluidId::new(fluid),
        Pressure::from_bar(p_bar),
        Enthalpy::from_joules_per_kg(h_kj_kg * 1000.0),
    );
    a.connect(b).expect("port connection ok").0
}

/// Creates screw compressor performance curves representing a ~200 kW screw
/// refrigerating unit at 50 Hz (R134a).
///
/// SST reference: +3°C = 276.15 K
/// SDT reference: +50°C = 323.15 K
fn make_screw_curves() -> ScrewPerformanceCurves {
    // Bilinear approximation:
    // ṁ_suc [kg/s] = 1.20 + 0.003×(SST-276) - 0.002×(SDT-323) + 1e-5×(SST-276)×(SDT-323)
    // W_shaft [W]  = 55000 + 200×(SST-276) - 300×(SDT-323) + 0.5×…
    ScrewPerformanceCurves::with_fixed_eco_fraction(
        Polynomial2D::bilinear(1.20, 0.003, -0.002, 0.000_01),
        Polynomial2D::bilinear(55_000.0, 200.0, -300.0, 0.5),
        0.12, // 12% economizer fraction
    )
}

// ─────────────────────────────────────────────────────────────────────────────
// Mock components used for sections not yet wired with real residuals
// (FloodedEvaporator, Drum, Economizer, ExpansionValve, BrineSource/Sink,
//  FlowSplitter/Merger — these already exist as real components, but for this
//  topology test we use mocks to isolate the new components under test)
// ─────────────────────────────────────────────────────────────────────────────

/// Generic mock component: all residuals = 0, n_equations configurable.
struct Mock {
    n: usize,
    #[allow(dead_code)] // Stored for fixture completeness; not asserted in this test.
    circuit_id: CircuitId,
}

impl Mock {
    fn new(n: usize, circuit: u16) -> Self {
        Self {
            n,
            circuit_id: CircuitId(circuit),
        }
    }
}

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

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

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

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

    fn port_mass_flows(&self, _state: &StateSlice) -> Result<Vec<MassFlow>, ComponentError> {
        Ok(vec![MassFlow::from_kg_per_s(1.0)])
    }

    fn energy_transfers(&self, _state: &StateSlice) -> Option<(Power, Power)> {
        Some((Power::from_watts(0.0), Power::from_watts(0.0)))
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 1: ScrewEconomizerCompressor topology
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_screw_compressor_creation_and_residuals() {
    let suc = make_port("R134a", 3.2, 400.0);
    let dis = make_port("R134a", 12.8, 440.0);
    let eco = make_port("R134a", 6.4, 260.0);

    let comp =
        ScrewEconomizerCompressor::new(make_screw_curves(), "R134a", 50.0, 0.92, suc, dis, eco)
            .expect("compressor creation ok");

    assert_eq!(comp.n_equations(), 6);

    // CM1.3: ṁ values are edge unknowns at indices 0,1,2; W_shaft is internal at index 3.
    // set_system_context(global_state_offset=3, [(suc:m=0,p,h), (eco:m=1,p,h), (dis:m=2,p,h)])
    let mut comp = comp;
    comp.set_system_context(3, &[(0, 0, 0), (1, 0, 0), (2, 0, 0)]);

    let state = vec![
        1.2,      // state[0] = ṁ_suc [kg/s]
        0.144,    // state[1] = ṁ_eco [kg/s] = 12% × 1.2
        1.344,    // state[2] = ṁ_dis [kg/s] = ṁ_suc + ṁ_eco
        55_000.0, // state[3] = W_shaft [W] (internal state at offset 3)
    ];
    let mut residuals = vec![0.0; 6];
    comp.compute_residuals(&state, &mut residuals)
        .expect("residuals computed");

    // All residuals must be finite
    for (i, r) in residuals.iter().enumerate() {
        assert!(r.is_finite(), "residual[{}] = {} not finite", i, r);
    }

    // residuals[5] (mass balance: ṁ_dis - ṁ_suc - ṁ_eco) should be ≈ 0
    assert!(
        residuals[5].abs() < 1e-10,
        "Mass balance residual: {}",
        residuals[5]
    );
    // Residual[4] (shaft power balance): W_calc - W_state
    println!("Screw residuals: {:?}", residuals);
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 2: VFD frequency scaling
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_screw_vfd_scaling() {
    let suc = make_port("R134a", 3.2, 400.0);
    let dis = make_port("R134a", 12.8, 440.0);
    let eco = make_port("R134a", 6.4, 260.0);

    let mut comp =
        ScrewEconomizerCompressor::new(make_screw_curves(), "R134a", 50.0, 0.92, suc, dis, eco)
            .unwrap();

    // CM1.3: set_system_context so ṁ indices are 0,1,2 and W_shaft is at index 3
    comp.set_system_context(3, &[(0, 0, 0), (1, 0, 0), (2, 0, 0)]);

    // At full speed (50 Hz): compute mass flow residual
    let state_full = vec![1.2, 0.144, 1.344, 55_000.0];
    let mut r_full = vec![0.0; 6];
    comp.compute_residuals(&state_full, &mut r_full).unwrap();
    let m_error_full = r_full[0].abs();

    // At 40 Hz (80%): mass flow should be ~80% of full speed
    comp.set_frequency_hz(40.0).unwrap();
    assert!((comp.frequency_ratio() - 0.8).abs() < 1e-10);

    let state_reduced = vec![0.96, 0.115, 1.075, 44_000.0];
    let mut r_reduced = vec![0.0; 6];
    comp.compute_residuals(&state_reduced, &mut r_reduced)
        .unwrap();
    let m_error_reduced = r_reduced[0].abs();

    println!(
        "VFD test: r[0] at 50Hz = {:.4}, at 40Hz = {:.4}",
        m_error_full, m_error_reduced
    );

    // Both should be finite
    assert!(m_error_full.is_finite());
    assert!(m_error_reduced.is_finite());
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 3: MCHX condenser coil UA correction
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_mchx_ua_correction_with_fan_speed() {
    // Coil bank: 4 coils, 15 kW/K each at design point (35°C, fan=100%)
    let ua_per_coil = 15_000.0; // W/K

    let mut coils: Vec<MchxCondenserCoil> = (0..4)
        .map(|i| MchxCondenserCoil::for_35c_ambient(ua_per_coil, i))
        .collect();

    // Total UA at full speed
    let ua_total_full: f64 = coils.iter().map(|c| c.ua_effective()).sum();
    assert!(
        (ua_total_full - 4.0 * ua_per_coil).abs() < 2000.0,
        "Total UA at full speed should be ≈ 60 kW/K, got {:.0}",
        ua_total_full
    );

    // Reduce fan 1 to 70% (anti-override scenario)
    coils[0].set_fan_speed_ratio(0.70);
    let ua_coil0_reduced = coils[0].ua_effective();
    let ua_coil0_full = coils[1].ua_effective(); // coil[1] still at 100%

    // UA at 70% speed = UA_nominal × 0.7^0.5 ≈ UA_nominal × 0.837
    let expected_ratio = 0.70_f64.sqrt();
    let actual_ratio = ua_coil0_reduced / ua_coil0_full;
    let tol = 0.02; // 2% tolerance
    assert!(
        (actual_ratio - expected_ratio).abs() < tol,
        "UA ratio expected {:.3}, got {:.3}",
        expected_ratio,
        actual_ratio
    );

    println!(
        "MCHX UA: full={:.0} W/K, at 70% fan={:.0} W/K (ratio={:.3})",
        ua_coil0_full, ua_coil0_reduced, actual_ratio
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 4: MCHX UA decreases at high ambient temperature
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_mchx_ua_ambient_temperature_effect() {
    let coil_35 = MchxCondenserCoil::for_35c_ambient(15_000.0, 0);
    let mut coil_45 = MchxCondenserCoil::for_35c_ambient(15_000.0, 0);

    coil_45.set_air_temperature_celsius(45.0);

    let ua_35 = coil_35.ua_effective();
    let ua_45 = coil_45.ua_effective();

    println!("UA at 35°C: {:.0} W/K, UA at 45°C: {:.0} W/K", ua_35, ua_45);

    // Higher ambient → lower air density → lower UA
    assert!(
        ua_45 < ua_35,
        "UA should decrease with higher ambient temperature"
    );

    // The reduction should be ~3% (density ratio: 1.12/1.09 ≈ 0.973)
    let density_35 = 1.12_f64;
    let density_45 = 101_325.0 / (287.058 * 318.15); // ≈ 1.109
    let expected_ratio = density_45 / density_35;
    let actual_ratio = ua_45 / ua_35;

    assert!(
        (actual_ratio - expected_ratio).abs() < 0.02,
        "Density ratio expected {:.4}, got {:.4}",
        expected_ratio,
        actual_ratio
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 5: 2-circuit system topology construction
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_two_circuit_chiller_topology() {
    let mut sys = System::new();

    // ── Circuit 0 (compressor + condenser side) ───────────────────────────────
    // Simplified topology using Mock components to validate graph construction:
    //
    // Screw comp → FlowSplitter → [CoilA, CoilB] → FlowMerger
    //                                                      → EXV → FloodedEvap
    //           ← Drum ← Economizer ←────────────────────────────┘

    // Screw compressor circuit 0
    let comp0_suc = make_port("R134a", 3.2, 400.0);
    let comp0_dis = make_port("R134a", 12.8, 440.0);
    let comp0_eco = make_port("R134a", 6.4, 260.0);
    let comp0 = ScrewEconomizerCompressor::new(
        make_screw_curves(),
        "R134a",
        50.0,
        0.92,
        comp0_suc,
        comp0_dis,
        comp0_eco,
    )
    .unwrap();

    let comp0_node = sys
        .add_component_to_circuit(Box::new(comp0), CircuitId::ZERO)
        .expect("add comp0");

    // 4 MCHX coils for circuit 0 (2 coils per circuit in this test)
    for i in 0..2 {
        let coil = MchxCondenserCoil::for_35c_ambient(15_000.0, i);
        let coil_node = sys
            .add_component_to_circuit(Box::new(coil), CircuitId::ZERO)
            .expect("add coil");
        sys.add_edge(comp0_node, coil_node).expect("comp→coil edge");
    }

    // FlowMerger (mock), EXV, FloodedEvap, Drum, Eco — all mock.
    // CM1.4: reduce mock equation counts so the topology stays square/under-constrained
    // (this test only validates construction, not a solve).
    let merger = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 0)), CircuitId::ZERO)
        .unwrap();
    let exv = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 0)), CircuitId::ZERO)
        .unwrap();
    let evap = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 0)), CircuitId::ZERO)
        .unwrap();
    let drum = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 0)), CircuitId::ZERO)
        .unwrap();
    let eco = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 0)), CircuitId::ZERO)
        .unwrap();

    // Connect: merger → exv → evap → drum → eco → comp0 (suction)
    sys.add_edge(merger, exv).unwrap();
    sys.add_edge(exv, evap).unwrap();
    sys.add_edge(evap, drum).unwrap();
    sys.add_edge(drum, eco).unwrap();
    sys.add_edge(eco, comp0_node).unwrap();
    sys.add_edge(comp0_node, merger).unwrap(); // closes loop via compressor

    // ── Circuit 1 (second independent compressor circuit) ─────────────────────
    let comp1_suc = make_port("R134a", 3.2, 400.0);
    let comp1_dis = make_port("R134a", 12.8, 440.0);
    let comp1_eco = make_port("R134a", 6.4, 260.0);
    let comp1 = ScrewEconomizerCompressor::new(
        make_screw_curves(),
        "R134a",
        50.0,
        0.92,
        comp1_suc,
        comp1_dis,
        comp1_eco,
    )
    .unwrap();

    let comp1_node = sys
        .add_component_to_circuit(Box::new(comp1), CircuitId(1))
        .expect("add comp1");

    // 2 coils for circuit 1
    for i in 2..4 {
        let coil = MchxCondenserCoil::for_35c_ambient(15_000.0, i);
        let coil_node = sys
            .add_component_to_circuit(Box::new(coil), CircuitId(1))
            .expect("add coil");
        sys.add_edge(comp1_node, coil_node)
            .expect("comp1→coil edge");
    }

    let merger1 = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 1)), CircuitId(1))
        .unwrap();
    let exv1 = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 1)), CircuitId(1))
        .unwrap();
    let evap1 = sys
        .add_component_to_circuit(Box::new(Mock::new(1, 1)), CircuitId(1))
        .unwrap();

    sys.add_edge(merger1, exv1).unwrap();
    sys.add_edge(exv1, evap1).unwrap();
    sys.add_edge(evap1, comp1_node).unwrap();
    sys.add_edge(comp1_node, merger1).unwrap();

    // ── Assert topology ───────────────────────────────────────────────────────
    assert_eq!(sys.circuit_count(), 2, "should have exactly 2 circuits");

    // Circuit 0: comp + 2 coils + merger + exv + evap + drum + eco = 9 nodes
    assert!(
        sys.circuit_nodes(CircuitId::ZERO).count() >= 8,
        "circuit 0 should have ≥8 nodes"
    );

    // Circuit 1: comp + 2 coils + merger + exv + evap = 6 nodes
    assert!(
        sys.circuit_nodes(CircuitId(1)).count() >= 5,
        "circuit 1 should have ≥5 nodes"
    );

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

    println!(
        "2-circuit chiller topology: {} nodes in circuit 0, {} in circuit 1",
        sys.circuit_nodes(CircuitId::ZERO).count(),
        sys.circuit_nodes(CircuitId(1)).count()
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 6: Fan anti-override control logic
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_fan_anti_override_speed_reduction() {
    // Simulate anti-override: when condensing pressure > limit,
    // reduce fan speed gradually until pressure stabilises.
    //
    // This test validates the MCHX UA response to fan speed changes,
    // which is the physical mechanism behind anti-override control.

    let ua_nominal = 15_000.0; // W/K per coil
    let mut coil = MchxCondenserCoil::for_35c_ambient(ua_nominal, 0);

    // Start at 100% fan speed
    assert!((coil.fan_speed_ratio() - 1.0).abs() < 1e-10);
    let ua_100 = coil.ua_effective();

    // Reduce to 80% (typical anti-override step)
    coil.set_fan_speed_ratio(0.80);
    let ua_80 = coil.ua_effective();

    // Reduce to 60%
    coil.set_fan_speed_ratio(0.60);
    let ua_60 = coil.ua_effective();

    // UA should decrease monotonically with fan speed
    assert!(ua_100 > ua_80, "UA should decrease from 100% to 80%");
    assert!(ua_80 > ua_60, "UA should decrease from 80% to 60%");

    // Reduction should follow power law: UA ∝ speed^0.5
    let ratio_80 = ua_80 / ua_100;
    let ratio_60 = ua_60 / ua_100;
    assert!(
        (ratio_80 - 0.80_f64.sqrt()).abs() < 0.03,
        "80% speed ratio: expected {:.3}, got {:.3}",
        0.80_f64.sqrt(),
        ratio_80
    );
    assert!(
        (ratio_60 - 0.60_f64.sqrt()).abs() < 0.03,
        "60% speed ratio: expected {:.3}, got {:.3}",
        0.60_f64.sqrt(),
        ratio_60
    );

    println!(
        "Anti-override UA: 100%={:.0}, 80%={:.0}, 60%={:.0} W/K",
        ua_100, ua_80, ua_60
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 7: Screw compressor off state — zero mass flow
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_screw_compressor_off_state_zero_flow() {
    let suc = make_port("R134a", 3.2, 400.0);
    let dis = make_port("R134a", 12.8, 440.0);
    let eco = make_port("R134a", 6.4, 260.0);

    let mut comp =
        ScrewEconomizerCompressor::new(make_screw_curves(), "R134a", 50.0, 0.92, suc, dis, eco)
            .unwrap();

    comp.set_state(OperationalState::Off).unwrap();
    // CM1.3: ṁ at indices 0,1,2; W_shaft at index 3
    comp.set_system_context(3, &[(0, 0, 0), (1, 0, 0), (2, 0, 0)]);

    let state = vec![0.0; 4];
    let mut residuals = vec![0.0; 6];
    comp.compute_residuals(&state, &mut residuals).unwrap();

    // In Off state: r[0]=ṁ_suc=0, r[1]=ṁ_eco=0, r[4]=W=0
    assert!(
        residuals[0].abs() < 1e-12,
        "Off: ṁ_suc residual should be 0"
    );
    assert!(
        residuals[1].abs() < 1e-12,
        "Off: ṁ_eco residual should be 0"
    );
    assert!(residuals[4].abs() < 1e-12, "Off: W residual should be 0");
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 8: 4-coil bank total capacity estimate
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_four_coil_bank_total_ua() {
    // Design: 4 coils, total UA = 60 kW/K, T_air=35°C
    // Expected: total condensing capacity ≈ 60 kW/K × (T_cond - T_air) ≈ 60 × 15 = 900 kW
    // (for T_cond = 50°C, ΔT_lm ≈ 15 K — rough estimate)

    let coils: Vec<MchxCondenserCoil> = (0..4)
        .map(|i| MchxCondenserCoil::for_35c_ambient(15_000.0, i))
        .collect();

    let total_ua: f64 = coils.iter().map(|c| c.ua_effective()).sum();

    println!(
        "4-coil bank total UA: {:.0} W/K = {:.1} kW/K",
        total_ua,
        total_ua / 1000.0
    );

    // Should be close to 60 kW/K (4 × 15 kW/K, with density ≈ 1 at design point)
    assert!(
        (total_ua - 60_000.0).abs() < 3_000.0,
        "Total UA should be ≈ 60 kW/K, got {:.1} kW/K",
        total_ua / 1000.0
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 9: Cross-circuit connection rejected
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_cross_circuit_connection_rejected() {
    let mut sys = System::new();

    let n0 = sys
        .add_component_to_circuit(Box::new(Mock::new(2, 0)), CircuitId::ZERO)
        .unwrap();
    let n1 = sys
        .add_component_to_circuit(Box::new(Mock::new(2, 1)), CircuitId(1))
        .unwrap();

    let result = sys.add_edge(n0, n1);
    assert!(
        matches!(result, Err(TopologyError::CrossCircuitConnection { .. })),
        "Cross-circuit edge should be rejected"
    );
}

// ─────────────────────────────────────────────────────────────────────────────
// Test 10: Screw compressor energy balance sanity check
// ─────────────────────────────────────────────────────────────────────────────

#[test]
fn test_screw_energy_balance() {
    let suc = make_port("R134a", 3.2, 400.0);
    let dis = make_port("R134a", 12.8, 440.0);
    let eco = make_port("R134a", 6.4, 260.0);

    let comp =
        ScrewEconomizerCompressor::new(make_screw_curves(), "R134a", 50.0, 0.92, suc, dis, eco)
            .unwrap();

    // At this operating point:
    // h_suc=400 kJ/kg, h_dis=440 kJ/kg, h_eco=260 kJ/kg
    // ṁ_suc=1.2 kg/s, ṁ_eco=0.144 kg/s, ṁ_total=1.344 kg/s
    // First law (fluid side): ṁ_suc×h_suc + ṁ_eco×h_eco + W_fluid = ṁ_total×h_dis
    // W_fluid = W_shaft × η_mech
    // W_shaft = (ΔH) / η_mech

    let m_suc = 1.2_f64;
    let m_eco = 0.144_f64;
    let m_total = m_suc + m_eco;
    let h_suc = 400_000.0_f64;
    let h_dis = 440_000.0_f64;
    let h_eco = 260_000.0_f64;
    let eta_mech = 0.92_f64;

    let delta_h = m_total * h_dis - m_suc * h_suc - m_eco * h_eco;
    let w_shaft = delta_h / eta_mech;
    let w_fluid = w_shaft * eta_mech; // == delta_h
    println!(
        "Shaft power: {:.0} W = {:.1} kW, Fluid power: {:.0} W",
        w_shaft,
        w_shaft / 1000.0,
        w_fluid
    );

    // Verify: W_shaft closes the energy balance via residual[2]
    // CM1.3 state layout: [m_suc, m_eco, m_dis, w_shaft] — enthalpies come from ports
    let mut comp = comp;
    comp.set_system_context(3, &[(0, 0, 0), (1, 0, 0), (2, 0, 0)]);
    let state = vec![m_suc, m_eco, m_suc + m_eco, w_shaft];
    let mut residuals = vec![0.0; 6];
    comp.compute_residuals(&state, &mut residuals).unwrap();

    // residual[2] = (ṁ_suc×h_suc + ṁ_eco×h_eco + W_shaft×η) - ṁ_total×h_dis
    println!("Energy balance residual: {:.4} J/s", residuals[2]);
    assert!(
        residuals[2].abs() < 1.0,
        "Energy balance residual should be < 1 W, got {:.4}",
        residuals[2]
    );
}
