//! Thermal load component: the cold-side receiver of an inter-circuit
//! thermal coupling.
//!
//! A `ThermalLoad` models a hydronic load segment (e.g. the cooling-water side
//! of a shared heat exchanger) that receives an **externally-determined heat
//! rate Q [W]** from the solver's thermal-coupling layer.
//!
//! Following the BOLT/Modelica boundary pattern, the loop's pressure and
//! inlet temperature are fixed by **boundary components** (`BrineSource` for
//! P_set/T_in, `BrineSink` for the back-pressure, temperature left free), NOT
//! by this component. The load itself only contributes:
//!
//! ```text
//! r0:  ṁ − ṁ_design                       (imposed design flow)
//! r1:  ṁ_design·(h_out − h_in) − Q_ext    (energy balance, Q_ext = state[q_idx])
//! ```
//!
//! Topology (mirrors `BOLT.BoundaryNode.Coolant.Source → HX → Sink`):
//!
//! ```text
//! BrineSource(P,T_in) ──edge──▶ ThermalLoad ──edge──▶ BrineSink(P_back, T free)
//! ```
//!
//! The energy balance uses the **design** flow (a constant), not the ṁ unknown:
//! r0 already pins ṁ = ṁ_design exactly, so the two forms coincide at the
//! solution — but the constant form keeps the block linear and avoids a
//! structurally singular Jacobian when the initializer starts at ṁ = 0
//! (the bilinear form's ∂r1/∂h_out = ṁ would then vanish).
//!
//! `Q_ext` is read from the per-coupling state unknown wired by
//! `System::finalize()` via [`Component::set_external_heat_index`]. When no
//! coupling is wired, `Q_ext = 0` and the load is an adiabatic pass-through
//! with imposed flow. The outlet temperature is **emergent**:
//! `T_out = T_in + Q/(ṁ·cp)`.

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

/// Cold-side thermal load with imposed design mass flow and externally-injected
/// heat rate (see module docs). Pair with `BrineSource`/`BrineSink` boundaries.
#[derive(Debug, Clone)]
pub struct ThermalLoad {
    name: String,
    /// Imposed design mass flow [kg/s].
    mass_flow_kg_s: f64,
    /// State index of the external heat unknown Q [W] (wired by the solver).
    q_idx: Option<usize>,
    /// Inlet edge (m, h) state indices.
    inlet_m_idx: Option<usize>,
    inlet_h_idx: Option<usize>,
    /// Outlet edge h state index.
    outlet_h_idx: Option<usize>,
    circuit_id: CircuitId,
    operational_state: OperationalState,
}

impl ThermalLoad {
    /// Creates a thermal load with the given imposed design mass flow [kg/s].
    pub fn new(name: impl Into<String>, mass_flow_kg_s: f64) -> Result<Self, ComponentError> {
        if !(mass_flow_kg_s.is_finite() && mass_flow_kg_s > 0.0) {
            return Err(ComponentError::InvalidState(
                "ThermalLoad: mass_flow_kg_s must be positive".to_string(),
            ));
        }
        Ok(Self {
            name: name.into(),
            mass_flow_kg_s,
            q_idx: None,
            inlet_m_idx: None,
            inlet_h_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 imposed design mass flow [kg/s].
    pub fn mass_flow_kg_s(&self) -> f64 {
        self.mass_flow_kg_s
    }

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

    /// External heat rate Q [W] read from the state (0 when not wired).
    fn q_ext(&self, state: &StateSlice) -> f64 {
        match self.q_idx {
            Some(i) if i < state.len() && state[i].is_finite() => state[i],
            _ => 0.0,
        }
    }

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

impl Component for ThermalLoad {
    fn set_system_context(
        &mut self,
        _state_offset: usize,
        external_edge_state_indices: &[(usize, usize, usize)],
    ) {
        // Layout: [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_h_idx = Some(h);
        }
        if external_edge_state_indices.len() >= 2 {
            let (_, _p, h) = external_edge_state_indices[1];
            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_h, out_h)) = self.wired_indices() else {
            return Err(ComponentError::InvalidState(
                "ThermalLoad requires live inlet/outlet edge state indices".to_string(),
            ));
        };
        let m_dot = state[m_idx];
        let h_in = state[in_h];
        let h_out = state[out_h];

        match self.operational_state {
            OperationalState::Off => {
                residuals[0] = m_dot; // ṁ = 0
                residuals[1] = h_out - h_in; // adiabatic
            }
            OperationalState::Bypass => {
                residuals[0] = m_dot - self.mass_flow_kg_s;
                residuals[1] = h_out - h_in; // adiabatic pass-through
            }
            OperationalState::On => {
                residuals[0] = m_dot - self.mass_flow_kg_s;
                // Linear form with the design flow (see module docs).
                residuals[1] = self.mass_flow_kg_s * (h_out - h_in) - self.q_ext(state);
            }
        }
        Ok(())
    }

    fn jacobian_entries(
        &self,
        _state: &StateSlice,
        jacobian: &mut JacobianBuilder,
    ) -> Result<(), ComponentError> {
        let Some((m_idx, in_h, out_h)) = self.wired_indices() else {
            return Ok(());
        };
        jacobian.add_entry(0, m_idx, 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 => {
                jacobian.add_entry(1, out_h, self.mass_flow_kg_s);
                jacobian.add_entry(1, in_h, -self.mass_flow_kg_s);
                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 m = self
            .inlet_m_idx
            .filter(|&i| i < state.len())
            .map(|i| state[i])
            .unwrap_or(0.0);
        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(
                "ThermalLoad: not wired to system context".to_string(),
            ));
        };
        if in_h >= state.len() || out_h >= state.len() {
            return Err(ComponentError::CalculationFailed(
                "ThermalLoad: state indices out of bounds".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)> {
        // Heat added TO the component (from the coupled hot circuit) is positive.
        Some((
            entropyk_core::Power::from_watts(self.q_ext(state)),
            entropyk_core::Power::from_watts(0.0),
        ))
    }

    fn counts_in_cycle_performance(&self) -> bool {
        // The absorbed Q is the primary cycle's rejected duty — do not count
        // it as additional cooling capacity in COP aggregation.
        false
    }

    fn measure_output(&self, kind: MeasuredOutput, state: &StateSlice) -> Option<f64> {
        match kind {
            // Received coupled heat [W].
            MeasuredOutput::Capacity | MeasuredOutput::HeatTransferRate => {
                Some(self.q_ext(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!(
            "ThermalLoad(name={}, m_design={:.4} kg/s, circuit={})",
            self.name, self.mass_flow_kg_s, self.circuit_id.0
        )
    }

    fn to_params(&self) -> crate::ComponentParams {
        crate::ComponentParams::new("ThermalLoad")
            .with_param("name", self.name.as_str())
            .with_param("massFlowKgS", self.mass_flow_kg_s)
            .with_param("circuitId", self.circuit_id.0)
    }
}

impl StateManageable for ThermalLoad {
    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::*;

    fn wired_load() -> ThermalLoad {
        let mut load = ThermalLoad::new("cw_load", 0.5).unwrap();
        // inlet edge: (m=0, p=1, h=2); outlet edge: (m=0, p=3, h=4); Q at 5.
        load.set_system_context(0, &[(0, 1, 2), (0, 3, 4)]);
        load.set_external_heat_index(5);
        load
    }

    #[test]
    fn test_rejects_nonpositive_flow() {
        assert!(ThermalLoad::new("x", 0.0).is_err());
        assert!(ThermalLoad::new("x", -1.0).is_err());
        assert!(ThermalLoad::new("x", f64::NAN).is_err());
    }

    #[test]
    fn test_energy_balance_residual_closed() {
        let load = wired_load();
        // ṁ = 0.5, h_in = 100 kJ/kg, h_out = 120 kJ/kg, Q = 10 kW
        let state = vec![0.5, 2.0e5, 100_000.0, 2.0e5, 120_000.0, 10_000.0];
        let mut r = vec![0.0; 2];
        load.compute_residuals(&state, &mut r).unwrap();
        assert!(r[0].abs() < 1e-12, "flow residual: {}", r[0]);
        assert!(r[1].abs() < 1e-9, "energy residual: {}", r[1]);
    }

    #[test]
    fn test_energy_residual_detects_imbalance() {
        let load = wired_load();
        // Q says 10 kW but the stream only picks up 5 kW.
        let state = vec![0.5, 2.0e5, 100_000.0, 2.0e5, 110_000.0, 10_000.0];
        let mut r = vec![0.0; 2];
        load.compute_residuals(&state, &mut r).unwrap();
        assert!(
            (r[1] - (-5_000.0)).abs() < 1e-9,
            "energy residual: {}",
            r[1]
        );
    }

    #[test]
    fn test_jacobian_matches_fd() {
        let load = wired_load();
        let state = vec![0.45, 2.0e5, 101_000.0, 1.9e5, 119_000.0, 9_500.0];
        let mut jac = JacobianBuilder::new();
        load.jacobian_entries(&state, &mut jac).unwrap();

        let eps = 1e-3;
        for col in [0usize, 2, 4, 5] {
            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]);
            load.compute_residuals(&sp, &mut rp).unwrap();
            load.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}"
                );
            }
        }
    }

    #[test]
    fn test_unwired_q_is_zero() {
        let mut load = ThermalLoad::new("cw_load", 0.5).unwrap();
        load.set_system_context(0, &[(0, 1, 2), (0, 3, 4)]);
        // No external heat wired: adiabatic (h_out must equal h_in).
        let state = vec![0.5, 2.0e5, 100_000.0, 2.0e5, 100_000.0];
        let mut r = vec![0.0; 2];
        load.compute_residuals(&state, &mut r).unwrap();
        assert!(r[1].abs() < 1e-12);
    }
}
