//! HeatSource — inline heat injection (BOLT `BoundaryNode.Heat.Source` equivalent).
//!
//! A `HeatSource` is a 2-port inline component that injects a heat rate Q [W]
//! into the stream flowing through it (Q < 0 extracts heat):
//!
//! ```text
//! r0:  P_out − P_in = 0                       (no hydraulic loss)
//! r1:  ṁ·(h_out − h_in) − Q_total = 0         (energy balance)
//! Q_total = q_fixed + Q_ext
//! ```
//!
//! Two heat sources compose:
//!
//! * **Fixed mode** (`q_fixed`, BOLT `Q_flow_fixed=true`): a constant duty,
//!   e.g. a known parasitic load or an electric heater.
//! * **Linked mode** (`Q_ext`, BOLT `use_Q_flow_in=true` with a
//!   `RealExpression` such as `Q = Comp_A.summary.Q_motorCooling`): the solver
//!   wires a per-coupling unknown via
//!   [`Component::set_external_heat_index`] whose value is closed against
//!   another component's measured output through a `thermal_couplings` entry
//!   (`hot_component` = duty provider, `cold_component` = this heat source).
//!
//! The Jacobian is exact (the bilinear term contributes
//! `∂r1/∂ṁ = h_out − h_in`, `∂r1/∂h = ±ṁ`, `∂r1/∂Q_ext = −1`).

use crate::state_machine::{CircuitId, OperationalState, StateManageable};
use crate::{
    Component, ComponentError, ConnectedPort, JacobianBuilder, MeasuredOutput, ResidualVector,
    StateSlice,
};

/// Inline heat injection with fixed and/or externally-linked duty (see module docs).
#[derive(Debug, Clone)]
pub struct HeatSource {
    name: String,
    /// Fixed heat rate [W] (positive = into the stream).
    q_fixed_w: f64,
    /// State index of the external heat unknown Q [W] (wired by the solver).
    q_idx: Option<usize>,
    inlet_m_idx: Option<usize>,
    inlet_p_idx: Option<usize>,
    inlet_h_idx: Option<usize>,
    outlet_p_idx: Option<usize>,
    outlet_h_idx: Option<usize>,
    circuit_id: CircuitId,
    operational_state: OperationalState,
}

impl HeatSource {
    /// Creates a heat source with a fixed duty `q_fixed_w` [W]
    /// (0.0 for a purely coupling-driven source).
    pub fn new(name: impl Into<String>, q_fixed_w: f64) -> Result<Self, ComponentError> {
        if !q_fixed_w.is_finite() {
            return Err(ComponentError::InvalidState(
                "HeatSource: q_fixed_w must be finite".to_string(),
            ));
        }
        Ok(Self {
            name: name.into(),
            q_fixed_w,
            q_idx: None,
            inlet_m_idx: None,
            inlet_p_idx: None,
            inlet_h_idx: None,
            outlet_p_idx: None,
            outlet_h_idx: None,
            circuit_id: CircuitId::default(),
            operational_state: OperationalState::default(),
        })
    }

    /// Returns the component name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns the fixed duty [W].
    pub fn q_fixed_w(&self) -> f64 {
        self.q_fixed_w
    }

    /// Returns the wired external-heat state index, if any.
    pub fn external_heat_index(&self) -> Option<usize> {
        self.q_idx
    }

    /// Total injected heat [W]: fixed + externally-linked.
    fn q_total(&self, state: &StateSlice) -> f64 {
        let q_ext = match self.q_idx {
            Some(i) if i < state.len() && state[i].is_finite() => state[i],
            _ => 0.0,
        };
        self.q_fixed_w + q_ext
    }

    fn wired(&self) -> Option<(usize, usize, usize, usize, usize)> {
        Some((
            self.inlet_m_idx?,
            self.inlet_p_idx?,
            self.inlet_h_idx?,
            self.outlet_p_idx?,
            self.outlet_h_idx?,
        ))
    }
}

impl Component for HeatSource {
    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize, usize)],
    ) {
        // [0] = incoming edge (inlet), [1] = outgoing edge (outlet).
        if !external_edge_state_indices.is_empty() {
            let (m, p, h) = external_edge_state_indices[0];
            self.inlet_m_idx = Some(m);
            self.inlet_p_idx = Some(p);
            self.inlet_h_idx = Some(h);
        }
        if external_edge_state_indices.len() >= 2 {
            let (_, p, h) = external_edge_state_indices[1];
            self.outlet_p_idx = Some(p);
            self.outlet_h_idx = Some(h);
        }
    }

    fn set_external_heat_index(&mut self, idx: usize) {
        self.q_idx = Some(idx);
    }

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

    fn compute_residuals(
        &self,
        state: &StateSlice,
        residuals: &mut ResidualVector,
    ) -> Result<(), ComponentError> {
        if residuals.len() < 2 {
            return Err(ComponentError::InvalidResidualDimensions {
                expected: 2,
                actual: residuals.len(),
            });
        }
        let Some((m_idx, in_p, in_h, out_p, out_h)) = self.wired() else {
            return Err(ComponentError::InvalidState(
                "HeatSource requires live inlet and outlet edge state indices".to_string(),
            ));
        };
        residuals[0] = state[out_p] - state[in_p];
        match self.operational_state {
            // Off / Bypass: adiabatic pass-through.
            OperationalState::Off | OperationalState::Bypass => {
                residuals[1] = state[out_h] - state[in_h];
            }
            OperationalState::On => {
                let m_dot = state[m_idx];
                residuals[1] = m_dot * (state[out_h] - state[in_h]) - self.q_total(state);
            }
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        let Some((m_idx, in_p, in_h, out_p, out_h)) = self.wired() else {
            return Err(ComponentError::InvalidState(
                "HeatSource Jacobian requires live inlet and outlet edge state indices".to_string(),
            ));
        };
        jacobian.add_entry(0, out_p, 1.0);
        jacobian.add_entry(0, in_p, -1.0);
        match self.operational_state {
            OperationalState::Off | OperationalState::Bypass => {
                jacobian.add_entry(1, out_h, 1.0);
                jacobian.add_entry(1, in_h, -1.0);
            }
            OperationalState::On => {
                let m_dot = state[m_idx];
                let dh = state[out_h] - state[in_h];
                if dh.abs() > 1e-14 {
                    jacobian.add_entry(1, m_idx, dh);
                }
                jacobian.add_entry(1, out_h, m_dot);
                jacobian.add_entry(1, in_h, -m_dot);
                if let Some(q_idx) = self.q_idx {
                    jacobian.add_entry(1, q_idx, -1.0);
                }
            }
        }
        Ok(())
    }

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

    fn port_mass_flows(
        &self,
        state: &StateSlice,
    ) -> Result<Vec<entropyk_core::MassFlow>, ComponentError> {
        let Some(m_idx) = self.inlet_m_idx.filter(|&i| i < state.len()) else {
            return Err(ComponentError::InvalidState(
                "HeatSource mass-flow reporting requires a live inlet mass-flow index".to_string(),
            ));
        };
        let m = state[m_idx];
        Ok(vec![
            entropyk_core::MassFlow::from_kg_per_s(m),
            entropyk_core::MassFlow::from_kg_per_s(-m),
        ])
    }

    fn port_enthalpies(
        &self,
        state: &StateSlice,
    ) -> Result<Vec<entropyk_core::Enthalpy>, ComponentError> {
        let (Some(in_h), Some(out_h)) = (self.inlet_h_idx, self.outlet_h_idx) else {
            return Err(ComponentError::CalculationFailed(
                "HeatSource: not wired to system context".to_string(),
            ));
        };
        Ok(vec![
            entropyk_core::Enthalpy::from_joules_per_kg(state[in_h]),
            entropyk_core::Enthalpy::from_joules_per_kg(state[out_h]),
        ])
    }

    fn energy_transfers(
        &self,
        state: &StateSlice,
    ) -> Option<(entropyk_core::Power, entropyk_core::Power)> {
        let q = match self.operational_state {
            OperationalState::On => self.q_total(state),
            _ => 0.0,
        };
        Some((
            entropyk_core::Power::from_watts(q),
            entropyk_core::Power::from_watts(0.0),
        ))
    }

    fn counts_in_cycle_performance(&self) -> bool {
        // Parasitic/auxiliary heat (e.g. motor cooling) must not inflate the
        // aggregated cooling capacity; it still enters First Law validation.
        false
    }

    fn measure_output(&self, kind: MeasuredOutput, state: &StateSlice) -> Option<f64> {
        match kind {
            MeasuredOutput::Capacity | MeasuredOutput::HeatTransferRate => {
                Some(self.q_total(state).abs())
            }
            MeasuredOutput::MassFlowRate => {
                let m_idx = self.inlet_m_idx?;
                (m_idx < state.len()).then(|| state[m_idx])
            }
            _ => None,
        }
    }

    fn signature(&self) -> String {
        format!(
            "HeatSource(name={}, q_fixed={:.1} W, circuit={})",
            self.name, self.q_fixed_w, self.circuit_id.0
        )
    }

    fn to_params(&self) -> crate::ComponentParams {
        crate::ComponentParams::new("HeatSource")
            .with_param("name", self.name.as_str())
            .with_param("qFixedW", self.q_fixed_w)
            .with_param("circuitId", self.circuit_id.0)
    }
}

impl StateManageable for HeatSource {
    fn state(&self) -> OperationalState {
        self.operational_state
    }

    fn set_state(&mut self, state: OperationalState) -> Result<(), ComponentError> {
        if self.operational_state.can_transition_to(state) {
            self.operational_state = state;
            Ok(())
        } else {
            Err(ComponentError::InvalidStateTransition {
                from: self.operational_state,
                to: state,
                reason: "Transition not allowed".to_string(),
            })
        }
    }

    fn can_transition_to(&self, target: OperationalState) -> bool {
        self.operational_state.can_transition_to(target)
    }

    fn circuit_id(&self) -> &CircuitId {
        &self.circuit_id
    }

    fn set_circuit_id(&mut self, circuit_id: CircuitId) {
        self.circuit_id = circuit_id;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Inlet edge (m=0, p=1, h=2), outlet edge (p=3, h=4), external Q at 5.
    fn wired(q_fixed: f64, external: bool) -> HeatSource {
        let mut hs = HeatSource::new("hs", q_fixed).unwrap();
        hs.set_system_context(0, &[(0, 1, 2), (0, 3, 4)]);
        if external {
            hs.set_external_heat_index(5);
        }
        hs
    }

    #[test]
    fn test_rejects_nonfinite_q() {
        assert!(HeatSource::new("x", f64::NAN).is_err());
        assert!(HeatSource::new("x", f64::INFINITY).is_err());
    }

    #[test]
    fn test_fixed_duty_energy_balance() {
        let hs = wired(5_000.0, false);
        // ṁ = 0.5, Δh = 10 kJ/kg ⇒ ṁ·Δh = 5 kW = q_fixed.
        let state = vec![0.5, 1.0e6, 400_000.0, 1.0e6, 410_000.0];
        let mut r = vec![0.0; 2];
        hs.compute_residuals(&state, &mut r).unwrap();
        assert!(r[0].abs() < 1e-12 && r[1].abs() < 1e-9, "{r:?}");
    }

    #[test]
    fn test_fixed_plus_external_compose() {
        let hs = wired(5_000.0, true);
        // Q_total = 5 kW + 3 kW ⇒ Δh = 16 kJ/kg at ṁ = 0.5.
        let state = vec![0.5, 1.0e6, 400_000.0, 1.0e6, 416_000.0, 3_000.0];
        let mut r = vec![0.0; 2];
        hs.compute_residuals(&state, &mut r).unwrap();
        assert!(r[1].abs() < 1e-9, "r1 = {}", r[1]);
    }

    #[test]
    fn test_negative_q_extracts_heat() {
        let hs = wired(-2_000.0, false);
        // Δh = −4 kJ/kg at ṁ = 0.5.
        let state = vec![0.5, 1.0e6, 400_000.0, 1.0e6, 396_000.0];
        let mut r = vec![0.0; 2];
        hs.compute_residuals(&state, &mut r).unwrap();
        assert!(r[1].abs() < 1e-9, "r1 = {}", r[1]);
    }

    #[test]
    fn test_jacobian_matches_fd() {
        let hs = wired(5_000.0, true);
        let state = vec![0.45, 1.0e6, 401_000.0, 0.98e6, 413_000.0, 2_500.0];
        let mut jac = JacobianBuilder::new();
        hs.jacobian_entries(&state, &mut jac).unwrap();

        let eps = 1e-3;
        for col in 0..6usize {
            let mut sp = state.clone();
            let mut sm = state.clone();
            sp[col] += eps;
            sm[col] -= eps;
            let (mut rp, mut rm) = (vec![0.0; 2], vec![0.0; 2]);
            hs.compute_residuals(&sp, &mut rp).unwrap();
            hs.compute_residuals(&sm, &mut rm).unwrap();
            for row in 0..2 {
                let fd = (rp[row] - rm[row]) / (2.0 * eps);
                let analytic: f64 = jac
                    .entries()
                    .iter()
                    .filter(|(r, c, _)| *r == row && *c == col)
                    .map(|(_, _, v)| *v)
                    .sum();
                assert!(
                    (fd - analytic).abs() < 1e-6 * (1.0 + fd.abs()),
                    "row {row} col {col}: fd {fd} vs analytic {analytic}"
                );
            }
        }
    }
}
