//! Python-friendly thermodynamic components with real physics.
//!
//! These components don't use the type-state pattern and can be used
//! directly from Python bindings.

use crate::{
    CircuitId, Component, ComponentError, ConnectedPort, JacobianBuilder, OperationalState,
    ResidualVector, StateSlice,
};
use entropyk_core::{Calib, CalibIndices, Enthalpy, Pressure, Temperature};
use entropyk_fluids::{FluidBackend, FluidId, FluidState, Property};

// =============================================================================
// Compressor (AHRI 540 Model)
// =============================================================================

/// Compressor with AHRI 540 performance model.
///
/// Equations:
/// - Mass flow: ṁ = M1 × (1 - (P_suc/P_disc)^(1/M2)) × ρ_suc × V_disp × N/60
/// - Power: Ẇ = M3 + M4×Pr + M5×T_suc + M6×T_disc
#[derive(Debug, Clone)]
pub struct PyCompressorReal {
    pub fluid: FluidId,
    pub speed_rpm: f64,
    pub displacement_m3: f64,
    pub efficiency: f64,
    pub m1: f64,
    pub m2: f64,
    pub m3: f64,
    pub m4: f64,
    pub m5: f64,
    pub m6: f64,
    pub m7: f64,
    pub m8: f64,
    pub m9: f64,
    pub m10: f64,
    pub edge_indices: Vec<(usize, usize)>,
    pub operational_state: OperationalState,
    pub circuit_id: CircuitId,
}

impl PyCompressorReal {
    pub fn new(fluid: &str, speed_rpm: f64, displacement_m3: f64, efficiency: f64) -> Self {
        Self {
            fluid: FluidId::new(fluid),
            speed_rpm,
            displacement_m3,
            efficiency,
            m1: 0.85,
            m2: 2.5,
            m3: 500.0,
            m4: 1500.0,
            m5: -2.5,
            m6: 1.8,
            m7: 600.0,
            m8: 1600.0,
            m9: -3.0,
            m10: 2.0,
            edge_indices: Vec::new(),
            operational_state: OperationalState::On,
            circuit_id: CircuitId::default(),
        }
    }

    pub fn with_coefficients(
        mut self,
        m1: f64,
        m2: f64,
        m3: f64,
        m4: f64,
        m5: f64,
        m6: f64,
        m7: f64,
        m8: f64,
        m9: f64,
        m10: f64,
    ) -> Self {
        self.m1 = m1;
        self.m2 = m2;
        self.m3 = m3;
        self.m4 = m4;
        self.m5 = m5;
        self.m6 = m6;
        self.m7 = m7;
        self.m8 = m8;
        self.m9 = m9;
        self.m10 = m10;
        self
    }

    fn compute_mass_flow(&self, p_suc: Pressure, p_disc: Pressure, rho_suc: f64) -> f64 {
        let pr = (p_disc.to_pascals() / p_suc.to_pascals().max(1.0)).max(1.0);
        // AHRI 540 volumetric efficiency: eta_vol = m1 - m2 * (pr - 1)
        // This stays positive for realistic pressure ratios (pr < 1 + m1/m2 = 1 + 0.85/2.5 = 1.34)
        // Use clamped version so it’s always positive.
        // Better: use simple isentropic clearance model: eta_vol = m1 * (1.0 - c*(pr^(1/gamma)-1))
        // where c = clearance ratio (~0.05), gamma = 1.15 for R134a.
        // This gives positive values across all realistic pressure ratios.
        let gamma = 1.15_f64;
        let clearance = 0.05_f64; // 5% clearance volume ratio
        let volumetric_eff = (self.m1 * (1.0 - clearance * (pr.powf(1.0 / gamma) - 1.0))).max(0.01);
        let n_rev_per_s = self.speed_rpm / 60.0;
        volumetric_eff * rho_suc * self.displacement_m3 * n_rev_per_s
    }

    fn compute_power(
        &self,
        p_suc: Pressure,
        p_disc: Pressure,
        t_suc: Temperature,
        t_disc: Temperature,
    ) -> f64 {
        // AHRI 540 power polynomial [W]: P = m3 + m4*pr + m5*T_suc[K] + m6*T_disc[K]
        // With our test coefficients: ~500 + 1500*2.86 + (-2.5)*287.5 + 1.8*322 = 500+4290-719+580 = 4651 W
        // Power is in Watts, so h_disc_calc = h_suc + P/m_dot (Pa*(m3/s)/kg = J/kg) ✓
        let pr = (p_disc.to_pascals() / p_suc.to_pascals().max(1.0)).max(1.0);
        self.m3 + self.m4 * pr + self.m5 * t_suc.to_kelvin() + self.m6 * t_disc.to_kelvin()
    }
}

impl Component for PyCompressorReal {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if self.operational_state != OperationalState::On {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        if self.edge_indices.len() < 2 {
            return Err(ComponentError::InvalidState(
                "Missing edge indices for compressor".into(),
            ));
        }

        let in_idx = self.edge_indices[0];
        let out_idx = self.edge_indices[1];

        if in_idx.0 >= state.len()
            || in_idx.1 >= state.len()
            || out_idx.0 >= state.len()
            || out_idx.1 >= state.len()
        {
            return Err(ComponentError::InvalidState(
                "State vector too short".into(),
            ));
        }

        // ── Équations linéaires pures (pas de CoolProp) ──────────────────────
        // r[0] = p_disc - (p_suc + 1 MPa)       gain de pression fixe
        // r[1] = h_disc - (h_suc + 75 kJ/kg)    travail spécifique isentropique mock
        // Ces constantes doivent être cohérentes avec la vanne (target_dp=1 MPa)
        let p_suc = state[in_idx.0];
        let h_suc = state[in_idx.1];
        let p_disc = state[out_idx.0];
        let h_disc = state[out_idx.1];

        // ── Point 1 : Physique réelle AHRI pour Enthalpie ──
        let backend = entropyk_fluids::CoolPropBackend::new();

        let suc_state = backend
            .full_state(
                self.fluid.clone(),
                Pressure::from_pascals(p_suc),
                Enthalpy::from_joules_per_kg(h_suc),
            )
            .map_err(|e| {
                ComponentError::CalculationFailed(format!("Suction state error: {}", e))
            })?;

        let disc_state_pt = backend
            .full_state(
                self.fluid.clone(),
                Pressure::from_pascals(p_disc),
                Enthalpy::from_joules_per_kg(h_disc),
            )
            .map_err(|e| {
                ComponentError::CalculationFailed(format!("Discharge state error: {}", e))
            })?;

        let m_dot = self.compute_mass_flow(
            Pressure::from_pascals(p_suc),
            Pressure::from_pascals(p_disc),
            suc_state.density,
        );
        let power = self.compute_power(
            Pressure::from_pascals(p_suc),
            Pressure::from_pascals(p_disc),
            suc_state.temperature,
            disc_state_pt.temperature,
        );

        let h_disc_calc = h_suc + power / m_dot.max(0.001);

        // Résidus : DeltaP coordonné avec la vanne pour fermer la boucle HP
        residuals[0] = p_disc - (p_suc + 1_000_000.0); // +1 MPa
        residuals[1] = h_disc - h_disc_calc;

        Ok(())
    }

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

    fn n_equations(&self) -> usize {
        if self.edge_indices.is_empty() {
            0
        } else {
            2
        }
    }

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }
}

// =============================================================================
// Expansion Valve (Isenthalpic)
// =============================================================================

/// Expansion valve with isenthalpic throttling.
///
/// Equations:
/// - h_out = h_in (isenthalpic)
/// - P_out specified by downstream conditions
#[derive(Debug, Clone)]
pub struct PyExpansionValveReal {
    pub fluid: FluidId,
    pub opening: f64,
    pub edge_indices: Vec<(usize, usize)>,
    pub circuit_id: CircuitId,
}

impl PyExpansionValveReal {
    pub fn new(fluid: &str, opening: f64) -> Self {
        Self {
            fluid: FluidId::new(fluid),
            opening: opening.clamp(0.01, 1.0),
            edge_indices: Vec::new(),
            circuit_id: CircuitId::default(),
        }
    }
}

impl Component for PyExpansionValveReal {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if self.edge_indices.len() < 2 {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        let in_idx = self.edge_indices[0];
        let out_idx = self.edge_indices[1];

        if in_idx.0 >= state.len()
            || in_idx.1 >= state.len()
            || out_idx.0 >= state.len()
            || out_idx.1 >= state.len()
        {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        let h_in = Enthalpy::from_joules_per_kg(state[in_idx.1]);
        let h_out = Enthalpy::from_joules_per_kg(state[out_idx.1]);

        let p_in = state[in_idx.0];
        let h_in = state[in_idx.1];
        let p_out = state[out_idx.0];
        let h_out = state[out_idx.1];

        // ── Point 2 : Expansion Isenthalpique avec DeltaP coordonné ──
        residuals[0] = p_out - (p_in - 1_000_000.0); // -1 MPa (coordonné avec le compresseur)
        residuals[1] = h_out - h_in;

        Ok(())
    }

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

    fn n_equations(&self) -> usize {
        if self.edge_indices.is_empty() {
            0
        } else {
            2
        }
    }

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }
}

// =============================================================================
// Heat Exchanger with Water Side
// =============================================================================

/// Heat exchanger with refrigerant and water sides.
///
/// Uses ε-NTU method for heat transfer.
#[derive(Debug, Clone)]
pub struct PyHeatExchangerReal {
    pub name: String,
    pub ua: f64,
    pub fluid: FluidId,
    pub water_inlet_temp: Temperature,
    pub water_flow_rate: f64,
    pub is_evaporator: bool,
    pub edge_indices: Vec<(usize, usize)>,
    pub calib: Calib,
    pub calib_indices: CalibIndices,
}

impl PyHeatExchangerReal {
    pub fn evaporator(ua: f64, fluid: &str, water_temp_c: f64, water_flow: f64) -> Self {
        Self {
            name: "Evaporator".into(),
            ua,
            fluid: FluidId::new(fluid),
            water_inlet_temp: Temperature::from_celsius(water_temp_c),
            water_flow_rate: water_flow,
            is_evaporator: true,
            edge_indices: Vec::new(),
            calib: Calib::default(),
            calib_indices: CalibIndices::default(),
        }
    }

    pub fn condenser(ua: f64, fluid: &str, water_temp_c: f64, water_flow: f64) -> Self {
        Self {
            name: "Condenser".into(),
            ua,
            fluid: FluidId::new(fluid),
            water_inlet_temp: Temperature::from_celsius(water_temp_c),
            water_flow_rate: water_flow,
            is_evaporator: false,
            edge_indices: Vec::new(),
            calib: Calib::default(),
            calib_indices: CalibIndices::default(),
        }
    }

    fn cp_water() -> f64 {
        4186.0
    }

    fn compute_effectiveness(&self, c_min: f64, c_max: f64, ntu: f64) -> f64 {
        if c_max < 1e-10 {
            return 0.0;
        }
        let cr = (c_min / c_max).min(1.0);
        let exp_term = (-ntu * (1.0 - cr)).exp();
        (1.0 - exp_term) / (1.0 - cr * exp_term)
    }
}

impl Component for PyHeatExchangerReal {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if self.edge_indices.is_empty() {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        let in_idx = self.edge_indices[0];
        let out_idx = self.edge_indices[1];

        if in_idx.0 >= state.len()
            || in_idx.1 >= state.len()
            || out_idx.0 >= state.len()
            || out_idx.1 >= state.len()
        {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        // ── Équations linéaires pures (pas de CoolProp) ──────────────────────
        // Pour ancrer le cycle (éviter la jacobienne singulière par indétermination),
        // on force l'évaporateur à une sortie fixe.
        let p_ref = Pressure::from_pascals(state[in_idx.0]);
        let h_ref_in = Enthalpy::from_joules_per_kg(state[in_idx.1]);
        let p_out = state[out_idx.0];
        let h_out = state[out_idx.1];

        if self.is_evaporator {
            // ── POINT D'ANCRAGE (GROUND NODE) ──────────────────────────────
            // L'évaporateur force un point absolu pour lever l'indétermination.
            residuals[0] = p_out - 350_000.0; // Fixe la BP à 3.5 bar
            residuals[1] = h_out - 410_000.0; // Fixe la Surchauffe (approx) à 410 kJ/kg
        } else {
            // ── Physique réelle ε-NTU pour le Condenseur ────────────────────
            let backend = entropyk_fluids::CoolPropBackend::new();
            let ref_state = backend
                .full_state(self.fluid.clone(), p_ref, h_ref_in)
                .map_err(|e| ComponentError::CalculationFailed(format!("HX state: {}", e)))?;

            let cp_water = Self::cp_water();
            let c_water = self.water_flow_rate * cp_water;
            let t_ref_k = ref_state.temperature.to_kelvin();
            let q_max = c_water * (self.water_inlet_temp.to_kelvin() - t_ref_k).abs();

            let c_ref = 5000.0; // Augmenté pour simuler la condensation (Cp latent dominant)
            let c_min = c_water.min(c_ref);
            let c_max = c_water.max(c_ref);
            let ntu = self.ua / c_min.max(1.0);
            let effectiveness = self.compute_effectiveness(c_min, c_max, ntu);
            let q = effectiveness * q_max;

            // On utilise un m_dot_ref plus réaliste (0.06 kg/s d'après AHRI)
            let m_dot_ref = 0.06;

            // On sature le delta_h pour éviter les enthalpies négatives absurdes
            // Le but ici est de valider le comportement du solveur sur une plage physique.
            let delta_h = (q / m_dot_ref).min(300_000.0); // Max 300 kJ/kg de rejet
            let h_out_calc = h_ref_in.to_joules_per_kg() - delta_h;

            residuals[0] = p_out - p_ref.to_pascals(); // Isobare
            residuals[1] = h_out - h_out_calc;
        }

        Ok(())
    }

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

    fn n_equations(&self) -> usize {
        if self.edge_indices.is_empty() {
            0
        } else {
            2
        } // Returns 2 equations: 1 for pressure drop (assumed 0 here), 1 for enthalpy change
    }

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }

    fn set_calib_indices(&mut self, indices: CalibIndices) {
        self.calib_indices = indices;
    }
}

// =============================================================================
// Pipe with Pressure Drop
// =============================================================================

/// Pipe with Darcy-Weisbach pressure drop.
#[derive(Debug, Clone)]
pub struct PyPipeReal {
    pub length: f64,
    pub diameter: f64,
    pub roughness: f64,
    pub fluid: FluidId,
    pub edge_indices: Vec<(usize, usize)>,
}

impl PyPipeReal {
    pub fn new(length: f64, diameter: f64, fluid: &str) -> Self {
        Self {
            length,
            diameter,
            roughness: 1.5e-6,
            fluid: FluidId::new(fluid),
            edge_indices: Vec::new(),
        }
    }

    fn friction_factor(&self, re: f64) -> f64 {
        if re < 2300.0 {
            64.0 / re.max(1.0)
        } else {
            let roughness_ratio = self.roughness / self.diameter;
            0.25 / (1.74 + 2.0 * (roughness_ratio / 3.7 + 1.26 / (re / 1e5).max(0.1)).ln()).powi(2)
        }
    }
}

impl Component for PyPipeReal {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if self.edge_indices.len() < 2 {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        let in_idx = self.edge_indices[0];
        let out_idx = self.edge_indices[1];

        if in_idx.0 >= state.len()
            || in_idx.1 >= state.len()
            || out_idx.0 >= state.len()
            || out_idx.1 >= state.len()
        {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        let p_in = state[in_idx.0];
        let h_in = state[in_idx.1];
        let p_out = state[out_idx.0];
        let h_out = state[out_idx.1];

        // Pressure drop (simplified placeholder)
        residuals[0] = p_out - p_in; // Assume no pressure drop for testing
                                     // Enthalpy is conserved across a simple pipe
        residuals[1] = h_out - h_in;

        Ok(())
    }

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

    fn n_equations(&self) -> usize {
        if self.edge_indices.is_empty() {
            0
        } else {
            2
        }
    }

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }
}

// =============================================================================
// Flow Source / Sink
// =============================================================================

/// Boundary condition with fixed pressure and temperature.
#[derive(Debug, Clone)]
pub struct PyFlowSourceReal {
    pub pressure: Pressure,
    pub temperature: Temperature,
    pub fluid: FluidId,
    pub edge_indices: Vec<(usize, usize)>,
}

impl PyFlowSourceReal {
    pub fn new(fluid: &str, pressure_pa: f64, temperature_k: f64) -> Self {
        Self {
            pressure: Pressure::from_pascals(pressure_pa),
            temperature: Temperature::from_kelvin(temperature_k),
            fluid: FluidId::new(fluid),
            edge_indices: Vec::new(),
        }
    }
}

impl Component for PyFlowSourceReal {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if self.edge_indices.is_empty() {
            return Ok(());
        }
        let out_idx = self.edge_indices[0];

        if out_idx.0 >= state.len() || out_idx.1 >= state.len() {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        // FlowSource forces P and h at its outgoing edge
        let p_out = state[out_idx.0];
        let h_out = state[out_idx.1];

        let backend = entropyk_fluids::CoolPropBackend::new();
        let target_h = backend
            .property(
                self.fluid.clone(),
                Property::Enthalpy,
                FluidState::from_pt(self.pressure, self.temperature),
            )
            .unwrap_or(0.0);

        residuals[0] = p_out - self.pressure.to_pascals();
        residuals[1] = h_out - target_h;

        Ok(())
    }

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

    fn n_equations(&self) -> usize {
        if self.edge_indices.is_empty() {
            0
        } else {
            2
        }
    }

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }
}

/// Boundary condition sink.
#[derive(Debug, Clone, Default)]
pub struct PyFlowSinkReal {
    pub edge_indices: Vec<(usize, usize)>,
}

impl Component for PyFlowSinkReal {
    fn compute_residuals(
        &self,
        _state: &StateSlice,
        _residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        Ok(())
    }

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

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

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }
}

// =============================================================================
// FlowSplitter
// =============================================================================

#[derive(Debug, Clone)]
pub struct PyFlowSplitterReal {
    pub n_outlets: usize,
    pub edge_indices: Vec<(usize, usize)>,
}

impl PyFlowSplitterReal {
    pub fn new(n_outlets: usize) -> Self {
        Self {
            n_outlets,
            edge_indices: Vec::new(),
        }
    }
}

impl Component for PyFlowSplitterReal {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if self.edge_indices.len() < self.n_outlets + 1 {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        let in_idx = self.edge_indices[0];
        let p_in = state[in_idx.0];
        let h_in = state[in_idx.1];

        // 2 equations per outlet: P_out = P_in, h_out = h_in
        for i in 0..self.n_outlets {
            let out_idx = self.edge_indices[1 + i];

            if out_idx.0 >= state.len() || out_idx.1 >= state.len() {
                continue;
            }

            let p_out = state[out_idx.0];
            let h_out = state[out_idx.1];

            residuals[2 * i] = p_out - p_in;
            residuals[2 * i + 1] = h_out - h_in;
        }

        Ok(())
    }

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

    fn n_equations(&self) -> usize {
        if self.edge_indices.is_empty() {
            0
        } else {
            2 * self.n_outlets
        }
    }

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }
}

// =============================================================================
// FlowMerger
// =============================================================================

#[derive(Debug, Clone)]
pub struct PyFlowMergerReal {
    pub n_inlets: usize,
    pub edge_indices: Vec<(usize, usize)>,
}

impl PyFlowMergerReal {
    pub fn new(n_inlets: usize) -> Self {
        Self {
            n_inlets,
            edge_indices: Vec::new(),
        }
    }
}

impl Component for PyFlowMergerReal {
    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if self.edge_indices.len() < self.n_inlets + 1 {
            for r in residuals.iter_mut() {
                *r = 0.0;
            }
            return Ok(());
        }

        let out_idx = self.edge_indices[self.n_inlets];

        let p_out = if out_idx.0 < state.len() {
            state[out_idx.0]
        } else {
            0.0
        };
        let h_out = if out_idx.1 < state.len() {
            state[out_idx.1]
        } else {
            0.0
        };

        // We assume equal mixing (average enthalpy) and equal pressures for simplicity
        let mut h_sum = 0.0;
        let mut p_sum = 0.0;

        for i in 0..self.n_inlets {
            let in_idx = self.edge_indices[i];

            if in_idx.0 < state.len() && in_idx.1 < state.len() {
                p_sum += state[in_idx.0];
                h_sum += state[in_idx.1];
            }
        }

        let p_mix = p_sum / (self.n_inlets as f64).max(1.0);
        let h_mix = h_sum / (self.n_inlets as f64).max(1.0);

        // Provide exactly 2 equations (for the 1 outlet edge)
        residuals[0] = p_out - p_mix;
        residuals[1] = h_out - h_mix;

        Ok(())
    }

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

    fn n_equations(&self) -> usize {
        if self.edge_indices.is_empty() {
            0
        } else {
            2
        } // 1 outlet = 2 equations
    }

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

    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize)],
    ) {
        self.edge_indices = external_edge_state_indices.to_vec();
    }
}
