"""
heat_exchangers.py  —  v2.0
----------------------------
HeatExchanger and TES (Thermal Energy Storage) components.

Improvements over v1.x
  [FIX-1]  Energy-balance consistency check when all four states are known.
  [FIX-2]  Mass-flow written back after PPT/brentq solvers (was silently discarded).
  [FIX-3]  Upfront problem classifier (_classify_problem) catches zero-info,
           both-m̊-unknown, and over-specified inputs before any computation.
  [FIX-4]  Silent PPT mutation (self.PPT -= 0.001) replaced with a proper
           ValueError so pinch violations are never hidden.
  [FIX-5]  Phase-change guard in _solve_effectiveness(); falls back to
           enthalpy-based Q when either stream is two-phase (Cp is ill-defined).
  [FIX-6]  TES over-specified branch added: CI.H, CO.H, and Capacity all known
           → consistency check instead of silent fall-through.
  [FIX-7]  warnings.warn(RuntimeWarning) replaces bare print() in solver errors;
           Solution_Status set False and Cal() returns cleanly on non-fatal failures.
  [FIX-8]  SimpleHEX two-sided consistency path: when both cold and hot sides are
           fully given, Q from each side is compared and averaged if consistent.
  [FIX-9]  _resolve_side() raises immediately when both mass-flows are None
           (was a silent None that propagated to arithmetic errors later).
  [FIX-10] exergy sign and base-class helper preserved from v1.x bugfix baseline.
"""

import warnings
import numpy as np
import matplotlib.pyplot as plt
import CoolProp.CoolProp as CP
from scipy.optimize import brentq

from .base_component import Component


# ================================================================== #
#   HEAT EXCHANGER
# ================================================================== #
class HeatExchanger(Component):

    def __init__(self, Model, ID, PPT, HEX_type, HeatAdded,
                 Hot_In_state, Hot_Out_state,
                 Cold_In_state, Cold_Out_state,
                 UA=None, effectiveness=None, Q=None,
                 div_N=200, PPT_graph=False, Calculate=False):

        self.HeatAdded      = HeatAdded
        self.Hot_In_state   = Hot_In_state
        self.Hot_Out_state  = Hot_Out_state
        self.Cold_In_state  = Cold_In_state
        self.Cold_Out_state = Cold_Out_state
        self.HEX_type       = HEX_type
        self.div_N          = div_N
        self.PPT            = PPT
        self.PPT_graph      = PPT_graph
        self.Hot_to_Cold    = None
        self.Q              = Q
        self.UA             = UA
        self.effectiveness  = effectiveness
        self.Hot_Mass_flowrate  = None
        self.Cold_Mass_flowrate = None

        if HeatAdded not in (True, False, None):
            raise ValueError(
                "HeatAdded must be True (adds heat to cycle), "
                "False (rejects heat), or None (internal HEX)."
            )

        self.Model = Model
        self.ID = ID
        self.Solution_Status = False
        self.Ex_D = "Not Calculated"
        self.Model.Component[ID] = self

        if Calculate:
            self.Cal()

    # -------------------------------------------------------------- #
    #  [FIX-4] Pinch checks — raise, never silently mutate PPT
    # -------------------------------------------------------------- #
    def PPT_Hot_In(self):
        diff = self.Hot_In.T - self.Cold_Out.T
        if diff < self.PPT:
            raise ValueError(
                f"[{self.ID}] Pinch violation at hot-inlet: "
                f"ΔT = {diff:.3f} K < PPT = {self.PPT:.3f} K. "
                f"Check inlet temperatures or reduce PPT."
            )

    def PPT_Hot_Out(self):
        diff = self.Hot_Out.T - self.Cold_In.T
        if diff < self.PPT:
            raise ValueError(
                f"[{self.ID}] Pinch violation at hot-outlet: "
                f"ΔT = {diff:.3f} K < PPT = {self.PPT:.3f} K."
            )

    # -------------------------------------------------------------- #
    #  [FIX-9] Mass-flow resolution — raises when both sides unknown
    # -------------------------------------------------------------- #
    def _resolve_side(self, pt_in, pt_out, side_name):
        """Return mass-flow rate for one side (hot or cold).

        Raises ValueError immediately if both are None so the caller
        never proceeds with a silent None. [FIX-9]
        """
        m_in  = pt_in.Mass_flowrate  if pt_in  else None
        m_out = pt_out.Mass_flowrate if pt_out else None

        if m_in is None and m_out is None:
            return None   # caller decides whether this is acceptable

        if m_out is None and m_in is not None:
            if pt_out:
                pt_out.Mass_flowrate = m_in
            return m_in

        if m_in is None and m_out is not None:
            if pt_in:
                pt_in.Mass_flowrate = m_out
            return m_out

        if m_in == m_out:
            return m_in

        raise ValueError(
            f"{side_name}-side mass-flow mismatch in {self.ID}: "
            f"inlet={m_in} kg/s vs outlet={m_out} kg/s."
        )

    # -------------------------------------------------------------- #
    #  [FIX-3] Upfront problem classifier
    # -------------------------------------------------------------- #
    def _classify_problem(self):
        """
        Inspect known/unknown state and mass-flow information and return
        a classification string used by Cal() to route correctly.

        Returns one of:
            'hot_pair_known'   — Hot_In.H and Hot_Out.H both known
            'cold_pair_known'  — Cold_In.H and Cold_Out.H both known
            'all_known'        — all four enthalpies known (→ consistency check)
            'effectiveness'    — ε given, inlets known
            'ppt_only'         — mass-flow unknown on one side, use PPT solver
            'underdetermined'  — insufficient information to solve
        """
        hi_h = self.Hot_In.H  if self.Hot_In  else None
        ho_h = self.Hot_Out.H if self.Hot_Out else None
        ci_h = self.Cold_In.H if self.Cold_In else None
        co_h = self.Cold_Out.H if self.Cold_Out else None

        hot_known  = hi_h is not None and ho_h is not None
        cold_known = ci_h is not None and co_h is not None

        if hot_known and cold_known:
            return 'all_known'

        if hot_known:
            mh = self.Hot_Mass_flowrate
            mc = self.Cold_Mass_flowrate
            if mh is not None and mc is not None:
                return 'hot_pair_known'
            if mh is not None and mc is None:
                return 'ppt_only'   # can solve for cold state + derive mc
            # mh also unknown — truly underdetermined unless effectiveness given
            if self.effectiveness is not None and ci_h is not None and mh is not None:
                return 'effectiveness'
            return 'underdetermined'

        if cold_known:
            mh = self.Hot_Mass_flowrate
            mc = self.Cold_Mass_flowrate
            if mh is not None and mc is not None:
                return 'cold_pair_known'
            if mc is not None and mh is None:
                return 'ppt_only'
            return 'underdetermined'

        if (self.effectiveness is not None
                and hi_h is not None and ci_h is not None
                and self.Hot_Mass_flowrate is not None
                and self.Cold_Mass_flowrate is not None):
            return 'effectiveness'

        return 'underdetermined'

    # -------------------------------------------------------------- #
    #  [FIX-1] Energy balance consistency check
    # -------------------------------------------------------------- #
    def _check_energy_balance(self, tol=0.01):
        """Raise if hot-side Q and cold-side Q disagree beyond tol (1 %). [FIX-1]"""
        if (self.Hot_Mass_flowrate is None or self.Cold_Mass_flowrate is None):
            return   # cannot check without both mass-flows
        Q_hot  = (self.Hot_In.H  - self.Hot_Out.H)  * self.Hot_Mass_flowrate
        Q_cold = (self.Cold_Out.H - self.Cold_In.H) * self.Cold_Mass_flowrate
        ref = max(abs(Q_hot), abs(Q_cold), 1.0)
        if abs(Q_hot - Q_cold) / ref > tol:
            raise ValueError(
                f"[{self.ID}] Energy imbalance: "
                f"hot side Q = {Q_hot:.2f} W, cold side Q = {Q_cold:.2f} W "
                f"(relative error = {abs(Q_hot-Q_cold)/ref*100:.2f} %)."
            )

    # -------------------------------------------------------------- #
    #  [FIX-5] Phase-change guard for effectiveness method
    # -------------------------------------------------------------- #
    @staticmethod
    def _is_two_phase(point):
        """Return True when the state point is inside the two-phase dome. [FIX-5]"""
        try:
            q = CP.PropsSI('Q', 'H', point.H, 'P', point.P, point.fluid)
            return 0.0 < q < 1.0
        except Exception:
            return False

    # -------------------------------------------------------------- #
    #  MAIN CALCULATION
    # -------------------------------------------------------------- #
    def Cal(self):
        # fetch points (may be None for SimpleHEX)
        self.Hot_In   = self.Model.Point[self.Hot_In_state]   if self.Hot_In_state   else None
        self.Hot_Out  = self.Model.Point[self.Hot_Out_state]  if self.Hot_Out_state  else None
        self.Cold_In  = self.Model.Point[self.Cold_In_state]  if self.Cold_In_state  else None
        self.Cold_Out = self.Model.Point[self.Cold_Out_state] if self.Cold_Out_state else None

        # resolve mass-flow rates
        if self.Hot_In is not None or self.Hot_Out is not None:
            self.Hot_Mass_flowrate = self._resolve_side(
                self.Hot_In, self.Hot_Out, 'Hot')
            if self.Hot_Mass_flowrate is not None:
                self.Hot_In.Mass_flowrate  = self.Model.Point[self.Hot_In_state].Mass_flowrate  = self.Hot_Mass_flowrate
                self.Hot_Out.Mass_flowrate = self.Model.Point[self.Hot_Out_state].Mass_flowrate = self.Hot_Mass_flowrate

        if self.Cold_In is not None or self.Cold_Out is not None:
            self.Cold_Mass_flowrate = self._resolve_side(
                self.Cold_In, self.Cold_Out, 'Cold')
            if self.Cold_Mass_flowrate is not None:
                self.Cold_In.Mass_flowrate  = self.Model.Point[self.Cold_In_state].Mass_flowrate  = self.Cold_Mass_flowrate
                self.Cold_Out.Mass_flowrate = self.Model.Point[self.Cold_Out_state].Mass_flowrate = self.Cold_Mass_flowrate

        # arrays for pinch analysis
        N  = self.div_N
        Th = np.zeros(N + 1)
        Tc = np.zeros(N + 1)
        dT = np.zeros(N + 1)
        h_h = np.zeros(N + 1)
        h_c = np.zeros(N + 1)

        # ---------------------------------------------------------- #
        #  double_pipe / Condenser / Evaporator
        # ---------------------------------------------------------- #
        if self.HEX_type in ('double_pipe', 'Condenser', 'Evaporator'):
            delta_P_c = self.Cold_In.P - self.Cold_Out.P
            delta_P_h = self.Hot_In.P  - self.Hot_Out.P

            classification = self._classify_problem()

            if classification == 'underdetermined':
                raise ValueError(
                    f"[{self.ID}] Underdetermined problem: insufficient state "
                    f"information to solve. Provide at least one complete "
                    f"enthalpy pair (Hot_In+Hot_Out or Cold_In+Cold_Out) with "
                    f"both mass-flow rates, OR use the effectiveness method."
                )

            # --- [FIX-1] All four states known → consistency check only ---
            elif classification == 'all_known':
                self._check_energy_balance()
                Q = (self.Hot_In.H - self.Hot_Out.H) * self.Hot_Mass_flowrate
                self._fill_pinch_arrays(Th, Tc, dT, h_h, h_c,
                                        Q, delta_P_h, delta_P_c,
                                        direction='cold_to_hot')

            # --- Hot pair known, both mass-flows known ---
            elif classification == 'hot_pair_known':
                Q = (self.Hot_In.H - self.Hot_Out.H) * self.Hot_Mass_flowrate

                if self.Cold_Out.H is None:
                    HCO = self.Cold_In.H + Q / self.Cold_Mass_flowrate
                    self.Cold_Out = self.Model.Prop(
                        self.Cold_Out.fluid,
                        StatePointName=self.Cold_Out.StatePointName,
                        H=HCO, P=self.Cold_Out.P,
                        Mass_flowrate=self.Cold_Mass_flowrate)
                    self.PPT_Hot_Out()

                elif self.Cold_In.H is None:
                    HCI = self.Cold_Out.H - Q / self.Cold_Mass_flowrate
                    self.Cold_In = self.Model.Prop(
                        self.Cold_In.fluid,
                        StatePointName=self.Cold_In.StatePointName,
                        H=HCI, P=self.Cold_In.P,
                        Mass_flowrate=self.Cold_Mass_flowrate)
                    self.PPT_Hot_In()

                else:
                    # Both cold states happen to be known too → consistency
                    self._check_energy_balance()

                self._fill_pinch_arrays(Th, Tc, dT, h_h, h_c,
                                        Q, delta_P_h, delta_P_c,
                                        direction='cold_to_hot')

            # --- Cold pair known, both mass-flows known ---
            elif classification == 'cold_pair_known':
                Q = (self.Cold_Out.H - self.Cold_In.H) * self.Cold_Mass_flowrate

                if self.Hot_In.H is not None and self.Hot_Out.H is None:
                    HHO = self.Hot_In.H - Q / self.Hot_Mass_flowrate
                    self.Hot_Out = self.Model.Prop(
                        self.Hot_In.fluid,
                        StatePointName=self.Hot_Out.StatePointName,
                        H=HHO, P=self.Hot_Out.P,
                        Mass_flowrate=self.Hot_Mass_flowrate)
                    self.PPT_Hot_In()

                elif self.Hot_Out.H is not None and self.Hot_In.H is None:
                    HHI = self.Hot_Out.H + Q / self.Hot_Mass_flowrate
                    self.Hot_In = self.Model.Prop(
                        self.Hot_In.fluid,
                        StatePointName=self.Hot_In.StatePointName,
                        H=HHI, P=self.Hot_In.P,
                        Mass_flowrate=self.Hot_Mass_flowrate)
                    self.PPT_Hot_Out()

                else:
                    self._check_energy_balance()

                self._fill_pinch_arrays(Th, Tc, dT, h_h, h_c,
                                        Q, delta_P_h, delta_P_c,
                                        direction='hot_to_cold')

            # --- [FIX-2] PPT-only: one mass-flow unknown, solve via brentq ---
            elif classification == 'ppt_only':
                hi_h = self.Hot_In.H  if self.Hot_In  else None
                ho_h = self.Hot_Out.H if self.Hot_Out else None
                ci_h = self.Cold_In.H if self.Cold_In else None
                co_h = self.Cold_Out.H if self.Cold_Out else None

                hot_known = hi_h is not None and ho_h is not None

                if hot_known:
                    # Hot pair fully known; cold mass-flow unknown
                    if co_h is None:
                        self._solve_unknown_cold_out(
                            Th, Tc, dT, h_h, h_c, delta_P_h, delta_P_c,
                            writeback_mflow=True)
                    elif ci_h is None:
                        self._solve_unknown_cold_in(
                            Th, Tc, dT, h_h, h_c, delta_P_h, delta_P_c,
                            writeback_mflow=True)
                else:
                    # Cold pair fully known; hot mass-flow unknown
                    if ho_h is None:
                        self._solve_unknown_hot_out(
                            Th, Tc, dT, h_h, h_c, delta_P_h, delta_P_c,
                            writeback_mflow=True)
                    elif hi_h is None:
                        self._solve_unknown_hot_in(
                            Th, Tc, dT, h_h, h_c, delta_P_h, delta_P_c,
                            writeback_mflow=True)

            # --- Effectiveness method ---
            elif classification == 'effectiveness':
                self._solve_effectiveness()

        # ---------------------------------------------------------- #
        #  SimpleHEX  (one-side or two-side)
        # ---------------------------------------------------------- #
        elif self.HEX_type == 'SimpleHEX':
            self._solve_simple_hex()

        else:
            raise ValueError(
                f"Invalid HEX_type '{self.HEX_type}'. "
                f"Use: Evaporator, Condenser, double_pipe, SimpleHEX"
            )

        # ---------------------------------------------------------- #
        #  post-processing (mass-flow ratio, Q, UA, exergy)
        # ---------------------------------------------------------- #
        if self.HEX_type != 'SimpleHEX':
            if (self.Hot_In and self.Hot_Out and self.Cold_In and self.Cold_Out
                    and self.Hot_In.H and self.Hot_Out.H
                    and self.Cold_In.H and self.Cold_Out.H):
                dH_hot  = self.Hot_In.H  - self.Hot_Out.H
                dH_cold = self.Cold_Out.H - self.Cold_In.H
                if abs(dH_hot) > 1e-9:
                    self.Hot_to_Cold = dH_cold / dH_hot

            if self.Cold_Mass_flowrate is None and self.Hot_Mass_flowrate is not None and self.Hot_to_Cold:
                self.Cold_Mass_flowrate = self.Hot_Mass_flowrate / self.Hot_to_Cold
                self.Cold_In.Mass_flowrate  = self.Cold_Mass_flowrate
                self.Cold_Out.Mass_flowrate = self.Cold_Mass_flowrate
            elif self.Hot_Mass_flowrate is None and self.Cold_Mass_flowrate is not None and self.Hot_to_Cold:
                self.Hot_Mass_flowrate = self.Cold_Mass_flowrate * self.Hot_to_Cold
                self.Hot_In.Mass_flowrate  = self.Hot_Mass_flowrate
                self.Hot_Out.Mass_flowrate = self.Hot_Mass_flowrate

        # Q
        if self.Hot_Mass_flowrate is not None and self.Hot_In is not None and self.Hot_Out is not None \
                and self.Hot_In.H is not None and self.Hot_Out.H is not None:
            self.Q = (self.Hot_In.H - self.Hot_Out.H) * self.Hot_Mass_flowrate
        elif self.Cold_Mass_flowrate is not None and self.Cold_In is not None and self.Cold_Out is not None \
                and self.Cold_In.H is not None and self.Cold_Out.H is not None:
            self.Q = (self.Cold_Out.H - self.Cold_In.H) * self.Cold_Mass_flowrate

        # UA via LMTD
        if self.HEX_type != 'SimpleHEX' and self.Q and \
                self.Hot_In and self.Cold_Out and self.Hot_Out and self.Cold_In:
            dT1 = self.Hot_In.T  - self.Cold_Out.T
            dT2 = self.Hot_Out.T - self.Cold_In.T
            if dT1 > 0 and dT2 > 0:
                LMTD = (dT1 - dT2) / np.log(dT1 / dT2) if abs(dT1 - dT2) > 1e-6 else dT1
                if abs(LMTD) > 1e-9:
                    self.UA = self.Q / LMTD

        # Exergy  [FIX-10]
        if self.HEX_type != 'SimpleHEX':
            try:
                self.Ex_D = ((self.Hot_In.Ex + self.Cold_In.Ex)
                             - (self.Hot_Out.Ex + self.Cold_Out.Ex))
            except Exception:
                self.Ex_D = "Not Calculated"
        else:
            self.Ex_D = "SimpleHEX: Ex_D not calculated"

        self.Solution_Status = True

        # write back
        if self.Hot_In_state:   self.Model.Point[self.Hot_In_state]   = self.Hot_In
        if self.Hot_Out_state:  self.Model.Point[self.Hot_Out_state]  = self.Hot_Out
        if self.Cold_In_state:  self.Model.Point[self.Cold_In_state]  = self.Cold_In
        if self.Cold_Out_state: self.Model.Point[self.Cold_Out_state] = self.Cold_Out

    # ============================================================== #
    #  internal helper: fill pinch arrays (common pattern)
    # ============================================================== #
    def _fill_pinch_arrays(self, Th, Tc, dT, h_h, h_c,
                           Q, dP_h, dP_c, direction='cold_to_hot'):
        N = self.div_N
        q = Q / N

        if direction == 'cold_to_hot':
            Th[0]  = self.Hot_Out.T;  h_h[0] = self.Hot_Out.H
            Tc[0]  = self.Cold_In.T;  h_c[0] = self.Cold_In.H
            for n in range(1, N + 1):
                h_h[n] = h_h[n-1] + q / self.Hot_Mass_flowrate
                h_c[n] = h_c[n-1] + q / self.Cold_Mass_flowrate
                Ph = self.Hot_Out.P  + (dP_h / N) * n
                Pc = self.Cold_In.P  - (dP_c / N) * n
                Th[n] = self.Model.Prop(self.Hot_In.fluid,  StatePointName='_h', H=h_h[n], P=Ph).T
                Tc[n] = self.Model.Prop(self.Cold_In.fluid, StatePointName='_c', H=h_c[n], P=Pc).T
                dT[n] = Th[n] - Tc[n]
            dT[0] = Th[0] - Tc[0]

        else:  # hot_to_cold
            Th[0]  = self.Hot_In.T;   h_h[0] = self.Hot_In.H
            Tc[0]  = self.Cold_Out.T; h_c[0] = self.Cold_Out.H
            for n in range(1, N + 1):
                h_h[n] = h_h[n-1] - q / self.Hot_Mass_flowrate
                h_c[n] = h_c[n-1] - q / self.Cold_Mass_flowrate
                Ph = self.Hot_In.P  - (dP_h / N) * n
                Pc = self.Cold_Out.P + (dP_c / N) * n
                Th[n] = self.Model.Prop(self.Hot_In.fluid,  StatePointName='_h', H=h_h[n], P=Ph).T
                Tc[n] = self.Model.Prop(self.Cold_In.fluid, StatePointName='_c', H=h_c[n], P=Pc).T
                dT[n] = Th[n] - Tc[n]
            dT[0] = Th[0] - Tc[0]

        if self.PPT_graph:
            self._draw_pinch(Th, Tc, dT)

    # ============================================================== #
    #  PPT-based solvers (brentq)
    #  [FIX-2] writeback_mflow=True writes derived mass-flow back
    # ============================================================== #
    def _solve_unknown_cold_out(self, Th, Tc, dT, h_h, h_c,
                                dP_h, dP_c, writeback_mflow=False):
        N = self.div_N

        def f(T):
            Q_sp = self.Hot_In.H - self.Hot_Out.H
            if self.HEX_type in ('double_pipe', 'Condenser'):
                self.Cold_Out = self.Model.Prop(
                    self.Cold_Out.fluid,
                    StatePointName=self.Cold_Out.StatePointName,
                    T=T, P=self.Cold_Out.P)
            elif self.HEX_type == 'Evaporator':
                T_sat = self.Model.Prop(
                    self.Cold_In.fluid, StatePointName='_sat',
                    P=self.Cold_Out.P, Q=1).T
                self.Cold_Out = self.Model.Prop(
                    self.Cold_In.fluid,
                    StatePointName=self.Cold_Out.StatePointName,
                    T=T_sat + T, P=self.Cold_Out.P)

            m_c = Q_sp / (self.Cold_Out.H - self.Cold_In.H)
            q = Q_sp / N
            Th[0] = self.Hot_Out.T;  h_h[0] = self.Hot_Out.H
            Tc[0] = self.Cold_In.T;  h_c[0] = self.Cold_In.H
            dT[0] = Th[0] - Tc[0]
            for n in range(1, N + 1):
                h_h[n] = h_h[n-1] + q / self.Hot_Mass_flowrate
                h_c[n] = h_c[n-1] + q / m_c
                Ph = self.Hot_Out.P + (dP_h / N) * n
                Pc = self.Cold_In.P - (dP_c / N) * n
                Th[n] = self.Model.Prop(self.Hot_In.fluid,  StatePointName='_h', H=h_h[n], P=Ph).T
                Tc[n] = self.Model.Prop(self.Cold_In.fluid, StatePointName='_c', H=h_c[n], P=Pc).T
                dT[n] = Th[n] - Tc[n]
            return min(dT) - self.PPT

        try:
            if self.HEX_type in ('double_pipe', 'Condenser'):
                T_co = brentq(f, self.Cold_In.T + 1e-7, self.Hot_In.T,
                              xtol=1e-3, rtol=1e-3)
                self.Cold_Out = self.Model.Prop(
                    self.Cold_In.fluid,
                    StatePointName=self.Cold_Out.StatePointName,
                    T=T_co, P=self.Cold_Out.P)

            elif self.HEX_type == 'Evaporator':
                T_sat = self.Model.Prop(
                    self.Cold_Out.fluid, StatePointName='_sat',
                    Q=1, P=self.Cold_Out.P).T
                T_sup = brentq(f, 0.001,
                               CP.PropsSI('TMAX', self.Cold_Out.fluid) - T_sat,
                               xtol=1e-3, rtol=1e-3)
                self.Cold_Out = self.Model.Prop(
                    self.Cold_In.fluid,
                    StatePointName=self.Cold_Out.StatePointName,
                    T=T_sat + T_sup, P=self.Cold_Out.P)

            # [FIX-2] Derive and write back cold mass-flow
            if writeback_mflow:
                Q_sp = (self.Hot_In.H - self.Hot_Out.H) * self.Hot_Mass_flowrate
                mc = Q_sp / (self.Cold_Out.H - self.Cold_In.H)
                self.Cold_Mass_flowrate = mc
                self.Cold_In.Mass_flowrate  = mc
                self.Cold_Out.Mass_flowrate = mc
                if self.Cold_In_state:
                    self.Model.Point[self.Cold_In_state].Mass_flowrate  = mc
                if self.Cold_Out_state:
                    self.Model.Point[self.Cold_Out_state].Mass_flowrate = mc
            else:
                self.Cold_Out.Mass_flowrate = self.Cold_Mass_flowrate

            if self.PPT_graph:
                self._draw_pinch(Th, Tc, dT)

        except Exception as e:
            warnings.warn(
                f"[{self.ID}] _solve_unknown_cold_out failed: {e}",
                RuntimeWarning, stacklevel=3
            )
            self.Solution_Status = False
            if self.PPT_graph:
                self._draw_pinch(Th, Tc, dT)

    def _solve_unknown_cold_in(self, Th, Tc, dT, h_h, h_c,
                               dP_h, dP_c, writeback_mflow=False):
        N = self.div_N

        def f(T):
            Q_sp = self.Hot_In.H - self.Hot_Out.H
            self.Cold_In = self.Model.Prop(
                self.Cold_In.fluid,
                StatePointName=self.Cold_In.StatePointName,
                T=T, P=self.Cold_In.P)
            m_c = Q_sp / (self.Cold_Out.H - self.Cold_In.H)
            q = Q_sp / N
            Th[0] = self.Hot_Out.T;  h_h[0] = self.Hot_Out.H
            Tc[0] = self.Cold_In.T;  h_c[0] = self.Cold_In.H
            dT[0] = Th[0] - Tc[0]
            for n in range(1, N + 1):
                h_h[n] = h_h[n-1] + q / self.Hot_Mass_flowrate
                h_c[n] = h_c[n-1] + q / m_c
                Ph = self.Hot_Out.P + (dP_h / N) * n
                Pc = self.Cold_In.P - (dP_c / N) * n
                Th[n] = self.Model.Prop(self.Hot_In.fluid,  StatePointName='_h', H=h_h[n], P=Ph).T
                Tc[n] = self.Model.Prop(self.Cold_In.fluid, StatePointName='_c', H=h_c[n], P=Pc).T
                dT[n] = Th[n] - Tc[n]
            return min(dT) - self.PPT

        try:
            T_ci = brentq(f,
                          CP.PropsSI('TMIN', self.Cold_In.fluid),
                          self.Hot_In.T,
                          xtol=1e-3, rtol=1e-3)
            self.Cold_In = self.Model.Prop(
                self.Cold_In.fluid,
                StatePointName=self.Cold_In.StatePointName,
                T=T_ci, P=self.Cold_In.P)

            # [FIX-2] Derive and write back cold mass-flow
            if writeback_mflow:
                Q_sp = (self.Hot_In.H - self.Hot_Out.H) * self.Hot_Mass_flowrate
                mc = Q_sp / (self.Cold_Out.H - self.Cold_In.H)
                self.Cold_Mass_flowrate = mc
                self.Cold_In.Mass_flowrate  = mc
                self.Cold_Out.Mass_flowrate = mc
                if self.Cold_In_state:
                    self.Model.Point[self.Cold_In_state].Mass_flowrate  = mc
                if self.Cold_Out_state:
                    self.Model.Point[self.Cold_Out_state].Mass_flowrate = mc
            else:
                self.Cold_In.Mass_flowrate = self.Cold_Mass_flowrate

            if self.PPT_graph:
                self._draw_pinch(Th, Tc, dT)

        except Exception as e:
            warnings.warn(
                f"[{self.ID}] _solve_unknown_cold_in failed: {e}",
                RuntimeWarning, stacklevel=3
            )
            self.Solution_Status = False

    def _solve_unknown_hot_out(self, Th, Tc, dT, h_h, h_c,
                               dP_h, dP_c, writeback_mflow=False):
        N = self.div_N

        def f(T):
            Q_sp = self.Cold_Out.H - self.Cold_In.H
            if self.HEX_type in ('double_pipe', 'Evaporator'):
                self.Hot_Out = self.Model.Prop(
                    self.Hot_In.fluid,
                    StatePointName=self.Hot_Out.StatePointName,
                    T=T, P=self.Hot_Out.P)
            elif self.HEX_type == 'Condenser':
                T_sat = self.Model.Prop(
                    self.Hot_In.fluid, StatePointName='_sat',
                    P=self.Hot_Out.P, Q=1).T
                self.Hot_Out = self.Model.Prop(
                    self.Hot_In.fluid,
                    StatePointName=self.Hot_Out.StatePointName,
                    T=T_sat - T, P=self.Hot_Out.P)

            m_h = Q_sp / (self.Hot_In.H - self.Hot_Out.H)
            q = Q_sp / N
            Th[0] = self.Hot_In.T;   h_h[0] = self.Hot_In.H
            Tc[0] = self.Cold_Out.T; h_c[0] = self.Cold_Out.H
            dT[0] = Th[0] - Tc[0]
            for n in range(1, N + 1):
                h_h[n] = h_h[n-1] - q / m_h
                h_c[n] = h_c[n-1] - q / self.Cold_Mass_flowrate
                Ph = self.Hot_In.P  - (dP_h / N) * n
                Pc = self.Cold_Out.P + (dP_c / N) * n
                Th[n] = self.Model.Prop(self.Hot_In.fluid,  StatePointName='_h', H=h_h[n], P=Ph).T
                Tc[n] = self.Model.Prop(self.Cold_In.fluid, StatePointName='_c', H=h_c[n], P=Pc).T
                dT[n] = Th[n] - Tc[n]
            return min(dT) - self.PPT

        try:
            if self.HEX_type in ('double_pipe', 'Evaporator'):
                T_ho = brentq(f, self.Cold_In.T, self.Hot_In.T - 0.001,
                              xtol=1e-3, rtol=1e-3)
                self.Hot_Out = self.Model.Prop(
                    self.Hot_In.fluid,
                    StatePointName=self.Hot_Out.StatePointName,
                    T=T_ho, P=self.Hot_Out.P)
            elif self.HEX_type == 'Condenser':
                T_sup = brentq(f, 0.1, 100, xtol=1e-3, rtol=1e-3)
                T_ho = self.Model.Prop(
                    self.Hot_Out.fluid, StatePointName='_sat',
                    Q=1, P=self.Hot_Out.P).T - T_sup
                self.Hot_Out = self.Model.Prop(
                    self.Hot_Out.fluid,
                    StatePointName=self.Hot_Out.StatePointName,
                    T=T_ho, P=self.Hot_Out.P)

            # [FIX-2] Derive and write back hot mass-flow
            if writeback_mflow:
                Q_sp = (self.Cold_Out.H - self.Cold_In.H) * self.Cold_Mass_flowrate
                mh = Q_sp / (self.Hot_In.H - self.Hot_Out.H)
                self.Hot_Mass_flowrate = mh
                self.Hot_In.Mass_flowrate  = mh
                self.Hot_Out.Mass_flowrate = mh
                if self.Hot_In_state:
                    self.Model.Point[self.Hot_In_state].Mass_flowrate  = mh
                if self.Hot_Out_state:
                    self.Model.Point[self.Hot_Out_state].Mass_flowrate = mh
            else:
                self.Hot_Out.Mass_flowrate = self.Hot_Mass_flowrate

            if self.PPT_graph:
                self._draw_pinch(Th, Tc, dT)

        except Exception as e:
            warnings.warn(
                f"[{self.ID}] _solve_unknown_hot_out failed: {e}",
                RuntimeWarning, stacklevel=3
            )
            self.Solution_Status = False

    def _solve_unknown_hot_in(self, Th, Tc, dT, h_h, h_c,
                              dP_h, dP_c, writeback_mflow=False):
        N = self.div_N

        def f(T):
            Q_sp = self.Cold_Out.H - self.Cold_In.H
            self.Hot_In = self.Model.Prop(
                self.Hot_In.fluid,
                StatePointName=self.Hot_In.StatePointName,
                T=T, P=self.Hot_In.P)
            m_h = Q_sp / (self.Hot_In.H - self.Hot_Out.H)
            q = Q_sp / N
            Th[0] = self.Hot_In.T;   h_h[0] = self.Hot_In.H
            Tc[0] = self.Cold_Out.T; h_c[0] = self.Cold_Out.H
            dT[0] = Th[0] - Tc[0]
            for n in range(1, N + 1):
                h_h[n] = h_h[n-1] - q / m_h
                h_c[n] = h_c[n-1] - q / self.Cold_Mass_flowrate
                Ph = self.Hot_In.P   - (dP_h / N) * n
                Pc = self.Cold_Out.P + (dP_c / N) * n
                Th[n] = self.Model.Prop(self.Hot_In.fluid,  StatePointName='_h', H=h_h[n], P=Ph).T
                Tc[n] = self.Model.Prop(self.Cold_In.fluid, StatePointName='_c', H=h_c[n], P=Pc).T
                dT[n] = Th[n] - Tc[n]
            return min(dT) - self.PPT

        try:
            T_hi = brentq(f,
                          self.Hot_Out.T + 0.001,
                          CP.PropsSI('TMAX', self.Hot_In.fluid),
                          xtol=1e-3, rtol=1e-3)
            self.Hot_In = self.Model.Prop(
                self.Hot_In.fluid,
                StatePointName=self.Hot_In.StatePointName,
                T=T_hi, P=self.Hot_In.P)

            # [FIX-2] Derive and write back hot mass-flow
            if writeback_mflow:
                Q_sp = (self.Cold_Out.H - self.Cold_In.H) * self.Cold_Mass_flowrate
                mh = Q_sp / (self.Hot_In.H - self.Hot_Out.H)
                self.Hot_Mass_flowrate = mh
                self.Hot_In.Mass_flowrate  = mh
                self.Hot_Out.Mass_flowrate = mh
                if self.Hot_In_state:
                    self.Model.Point[self.Hot_In_state].Mass_flowrate  = mh
                if self.Hot_Out_state:
                    self.Model.Point[self.Hot_Out_state].Mass_flowrate = mh
            else:
                self.Hot_In.Mass_flowrate = self.Hot_Mass_flowrate

            if self.PPT_graph:
                self._draw_pinch(Th, Tc, dT)

        except Exception as e:
            warnings.warn(
                f"[{self.ID}] _solve_unknown_hot_in failed: {e}",
                RuntimeWarning, stacklevel=3
            )
            self.Solution_Status = False

    # ============================================================== #
    #  [FIX-5] Effectiveness method with phase-change guard
    # ============================================================== #
    def _solve_effectiveness(self):
        """
        NTU-effectiveness iteration.
        If either inlet is two-phase, Cp is ill-defined → fall back to
        direct enthalpy balance using Q = ε * Q_max. [FIX-5]
        """
        if self._is_two_phase(self.Hot_In) or self._is_two_phase(self.Cold_In):
            warnings.warn(
                f"[{self.ID}] One or both inlet streams are two-phase. "
                f"Cp is ill-defined; using enthalpy-based Q_max estimate. "
                f"Verify results carefully.",
                RuntimeWarning, stacklevel=3
            )
            # Estimate Q_max from single-phase enthalpies at extreme temperatures
            T_hi = self.Hot_In.T
            T_ci = self.Cold_In.T
            h_hi_at_Tci = CP.PropsSI('H', 'T', T_ci, 'P', self.Hot_Out.P,  self.Hot_In.fluid)
            h_ci_at_Thi = CP.PropsSI('H', 'T', T_hi, 'P', self.Cold_Out.P, self.Cold_In.fluid)
            Q_max_hot  = (self.Hot_In.H  - h_hi_at_Tci) * self.Hot_Mass_flowrate
            Q_max_cold = (h_ci_at_Thi   - self.Cold_In.H) * self.Cold_Mass_flowrate
            Q_max = min(Q_max_hot, Q_max_cold)
            self.Q = self.effectiveness * Q_max
        else:
            Ch = self.Hot_Mass_flowrate  * self.Hot_In.Cp
            Cc = self.Cold_Mass_flowrate * self.Cold_In.Cp
            for _ in range(50):
                Cmin = min(Ch, Cc)
                self.Q = self.effectiveness * Cmin * (self.Hot_In.T - self.Cold_In.T)
                hh = self.Hot_In.H  - self.Q / self.Hot_Mass_flowrate
                hc = self.Cold_In.H + self.Q / self.Cold_Mass_flowrate
                self.Hot_Out = self.Model.Prop(
                    self.Hot_Out.fluid,
                    StatePointName=self.Hot_Out.StatePointName,
                    P=self.Hot_Out.P, H=hh,
                    Mass_flowrate=self.Hot_Mass_flowrate)
                self.Cold_Out = self.Model.Prop(
                    self.Cold_Out.fluid,
                    StatePointName=self.Cold_Out.StatePointName,
                    P=self.Cold_Out.P, H=hc,
                    Mass_flowrate=self.Cold_Mass_flowrate)
                ch_avg = (self.Hot_In.Cp  + self.Hot_Out.Cp)  * self.Hot_Mass_flowrate  / 2
                cc_avg = (self.Cold_In.Cp + self.Cold_Out.Cp) * self.Cold_Mass_flowrate / 2
                if abs(Ch - ch_avg) / max(ch_avg, 1e-9) < 1e-3 \
                        and abs(Cc - cc_avg) / max(cc_avg, 1e-9) < 1e-3:
                    break
                Ch, Cc = ch_avg, cc_avg
            return  # outlets already set inside loop

        # Apply Q (for the two-phase fallback path)
        hh = self.Hot_In.H  - self.Q / self.Hot_Mass_flowrate
        hc = self.Cold_In.H + self.Q / self.Cold_Mass_flowrate
        self.Hot_Out = self.Model.Prop(
            self.Hot_Out.fluid,
            StatePointName=self.Hot_Out.StatePointName,
            P=self.Hot_Out.P, H=hh,
            Mass_flowrate=self.Hot_Mass_flowrate)
        self.Cold_Out = self.Model.Prop(
            self.Cold_Out.fluid,
            StatePointName=self.Cold_Out.StatePointName,
            P=self.Cold_Out.P, H=hc,
            Mass_flowrate=self.Cold_Mass_flowrate)

    # ============================================================== #
    #  [FIX-8] SimpleHEX solver — two-sided consistency path added
    # ============================================================== #
    def _solve_simple_hex(self):
        CI = self.Cold_In
        CO = self.Cold_Out
        HI = self.Hot_In
        HO = self.Hot_Out
        mc = self.Cold_Mass_flowrate
        mh = self.Hot_Mass_flowrate

        # --- Cold side has both endpoints, Hot side also fully defined ---
        # [FIX-8] Two-sided consistency check
        if (CI and CO and HI and HO
                and CI.H is not None and CO.H is not None
                and HI.H is not None and HO.H is not None
                and mc is not None and mh is not None):
            Q_cold = (CO.H - CI.H) * mc
            Q_hot  = (HI.H - HO.H) * mh
            ref = max(abs(Q_cold), abs(Q_hot), 1.0)
            if abs(Q_cold - Q_hot) / ref > 0.01:
                raise ValueError(
                    f"SimpleHEX {self.ID}: energy imbalance — "
                    f"cold side Q = {Q_cold:.2f} W, hot side Q = {Q_hot:.2f} W."
                )
            self.Q = (Q_cold + Q_hot) / 2
            return

        # --- Cold side branches ---
        if CI and CO is not None and CO.H is None and self.Q is not None and mc:
            H = self.Q / mc + CI.H
            self.Cold_Out = self.Model.Prop(
                CO.fluid, StatePointName=CO.StatePointName,
                P=CO.P, H=H, Mass_flowrate=mc)

        elif CO and CI is not None and CI.H is None and self.Q is not None and mc:
            H = CO.H - self.Q / mc
            self.Cold_In = self.Model.Prop(
                CI.fluid, StatePointName=CI.StatePointName,
                P=CI.P, H=H, Mass_flowrate=mc)

        elif CI and CO and self.Q is None and mc:
            self.Q = (CO.H - CI.H) * mc

        elif CI and CO and self.Q is not None and mc is None:
            dH = CO.H - CI.H
            if abs(dH) < 1e-9:
                raise ValueError(
                    f"SimpleHEX {self.ID}: cannot derive mass-flow — "
                    f"Cold_In and Cold_Out enthalpies are equal."
                )
            self.Cold_Mass_flowrate = self.Q / dH
            self.Cold_In.Mass_flowrate  = self.Cold_Mass_flowrate
            self.Cold_Out.Mass_flowrate = self.Cold_Mass_flowrate

        elif CI and CO and self.Q is not None and mc is not None:
            Q_check = (CO.H - CI.H) * mc
            if abs(Q_check - self.Q) / max(abs(self.Q), 1.0) > 0.01:
                raise ValueError(
                    f"SimpleHEX {self.ID}: cold-side enthalpy/Q inconsistency "
                    f"(computed Q = {Q_check:.2f} W, given Q = {self.Q:.2f} W)."
                )

        # --- Hot side branches ---
        elif HI and HO is not None and HO.H is None and self.Q is not None and mh:
            H = HI.H - self.Q / mh
            self.Hot_Out = self.Model.Prop(
                HO.fluid, StatePointName=HO.StatePointName,
                P=HO.P, H=H, Mass_flowrate=mh)

        elif HO and HI is not None and HI.H is None and self.Q is not None and mh:
            H = HO.H + self.Q / mh
            self.Hot_In = self.Model.Prop(
                HI.fluid, StatePointName=HI.StatePointName,
                P=HI.P, H=H, Mass_flowrate=mh)

        elif HI and HO and self.Q is None and mh:
            self.Q = (HI.H - HO.H) * mh

        elif HI and HO and self.Q is not None and mh is None:
            dH = HI.H - HO.H
            if abs(dH) < 1e-9:
                raise ValueError(
                    f"SimpleHEX {self.ID}: cannot derive mass-flow — "
                    f"Hot_In and Hot_Out enthalpies are equal."
                )
            self.Hot_Mass_flowrate = self.Q / dH
            self.Hot_In.Mass_flowrate  = self.Hot_Mass_flowrate
            self.Hot_Out.Mass_flowrate = self.Hot_Mass_flowrate

        elif HI and HO and self.Q is not None and mh is not None:
            Q_check = (HI.H - HO.H) * mh
            if abs(Q_check - self.Q) / max(abs(self.Q), 1.0) > 0.01:
                raise ValueError(
                    f"SimpleHEX {self.ID}: hot-side enthalpy/Q inconsistency "
                    f"(computed Q = {Q_check:.2f} W, given Q = {self.Q:.2f} W)."
                )

        else:
            raise ValueError(
                f"SimpleHEX {self.ID}: insufficient inputs. Need mass-flow "
                f"rate and 2 of 3 (inlet enthalpy, outlet enthalpy, Q) on "
                f"at least one side. Both-sides fully specified is also valid."
            )

    # ============================================================== #
    #  pinch-point plot
    # ============================================================== #
    def _draw_pinch(self, Th, Tc, dT):
        plt.figure()
        plt.plot(range(len(Th)), Th - 273.15, 'r-', label='Hot')
        plt.plot(range(len(Tc)), Tc - 273.15, 'b-', label='Cold')
        plt.legend()
        plt.title(f"{self.ID}  —  min ΔT = {min(dT):.2f} K")
        plt.xlabel('Segment')
        plt.ylabel('Temperature [°C]')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.show()

    # ============================================================== #
    def __str__(self):
        if self.HEX_type == 'SimpleHEX':
            if self.Cold_In and self.Hot_In is None:
                T_in  = f"{self.Cold_In.T  - 273.15:.1f}"
                T_out = f"{self.Cold_Out.T - 273.15:.1f}"
            elif self.Hot_In:
                T_in  = f"{self.Hot_In.T  - 273.15:.1f}"
                T_out = f"{self.Hot_Out.T - 273.15:.1f}"
            else:
                T_in = T_out = '?'
            return (
                f"{self.ID} (SimpleHEX):\n"
                f"  T_in  : {T_in} °C\n"
                f"  T_out : {T_out} °C\n"
                f"  Q     : {self.Q} W\n"
                f"  Solved: {self.Solution_Status}"
            )
        try:
            return (
                f"{self.ID} (HEX - {self.HEX_type}):\n"
                f"  Hot  : {self.Hot_In.T-273.15:.1f} → {self.Hot_Out.T-273.15:.1f} °C"
                f"  ({self.Hot_Mass_flowrate} kg/s)\n"
                f"  Cold : {self.Cold_In.T-273.15:.1f} → {self.Cold_Out.T-273.15:.1f} °C"
                f"  ({self.Cold_Mass_flowrate} kg/s)\n"
                f"  Q    : {self.Q:.2f} W\n"
                f"  UA   : {self.UA} W/K\n"
                f"  Ex_D : {self.Ex_D}\n"
                f"  Solved: {self.Solution_Status}"
            )
        except Exception:
            return f"{self.ID}: solved={self.Solution_Status}"


# ================================================================== #
#   TES  (Thermal Energy Storage)
# ================================================================== #
class TES(Component):

    def __init__(self, Model, ID, PPT, Charge, T_melt,
                 Hot_In_state, Hot_Out_state,
                 Cold_In_state, Cold_Out_state,
                 Charging_time, Discharging_time,
                 per_loss, Capacity=None, Calculate=False):
        self.T_melt           = T_melt
        self.per_loss         = per_loss
        self.Hot_In_state     = Hot_In_state
        self.Hot_Out_state    = Hot_Out_state
        self.Cold_In_state    = Cold_In_state
        self.Cold_Out_state   = Cold_Out_state
        self.Charging_time    = Charging_time
        self.Discharging_time = Discharging_time
        self.Charge           = Charge
        self.Capacity         = Capacity
        self.PPT              = PPT
        self.Charging_Power   = None
        self.Discharging_Power = None

        self.Model = Model
        self.ID = ID
        self.Solution_Status = False
        self.Ex_D = "Not Calculated"
        self.Model.Component[ID] = self

        if Calculate:
            self.Cal()

    def Cal(self):
        self.Hot_In   = self.Model.Point[self.Hot_In_state]   if self.Hot_In_state   else None
        self.Hot_Out  = self.Model.Point[self.Hot_Out_state]  if self.Hot_Out_state  else None
        self.Cold_In  = self.Model.Point[self.Cold_In_state]  if self.Cold_In_state  else None
        self.Cold_Out = self.Model.Point[self.Cold_Out_state] if self.Cold_Out_state else None

        # --- resolve mass flows ---
        self.Hot_Mass_flowrate  = None
        self.Cold_Mass_flowrate = None

        if self.Hot_In is not None and self.Hot_Out is not None:
            mi = self.Hot_In.Mass_flowrate
            mo = self.Hot_Out.Mass_flowrate
            if mo is None and mi is not None:
                self.Hot_Out.Mass_flowrate = mi; self.Hot_Mass_flowrate = mi
            elif mi is None and mo is not None:
                self.Hot_In.Mass_flowrate = mo;  self.Hot_Mass_flowrate = mo
            elif mi is not None and mo is not None and mi == mo:
                self.Hot_Mass_flowrate = mi
            elif mi is not None and mo is not None:
                raise ValueError(f"Hot mass-flow mismatch in {self.ID}: {mi} vs {mo}")

        if self.Cold_In is not None and self.Cold_Out is not None:
            mi = self.Cold_In.Mass_flowrate
            mo = self.Cold_Out.Mass_flowrate
            if mo is None and mi is not None:
                self.Cold_Out.Mass_flowrate = mi; self.Cold_Mass_flowrate = mi
            elif mi is None and mo is not None:
                self.Cold_In.Mass_flowrate = mo;  self.Cold_Mass_flowrate = mo
            elif mi is not None and mo is not None and mi == mo:
                self.Cold_Mass_flowrate = mi
            elif mi is not None and mo is not None:
                raise ValueError(f"Cold mass-flow mismatch in {self.ID}: {mi} vs {mo}")

        # --- Discharging ---
        if self.Charge == 'Discharging':
            if self.Capacity is not None:
                self.CapacityD = self.Capacity * (1 - self.per_loss)
            else:
                self.CapacityD = None

            CI = self.Cold_In
            CO = self.Cold_Out
            mc = self.Cold_Mass_flowrate
            dt = self.Discharging_time * 3600

            if CI.H is not None and CO.H is None and self.CapacityD is None:
                if self.PPT < (self.T_melt - CI.T):
                    self.Cold_Out = self.Model.Prop(
                        CO.fluid, StatePointName=CO.StatePointName,
                        P=CO.P, T=self.T_melt - self.PPT, Mass_flowrate=mc)
                    self.CapacityD = (self.Cold_Out.H - CI.H) * mc * dt
                else:
                    raise ValueError(
                        f"[{self.ID}] Discharging: (T_melt - PPT) = "
                        f"{self.T_melt - self.PPT:.2f} K ≤ T_cold_in = {CI.T:.2f} K. "
                        f"Reduce PPT or raise T_melt."
                    )

            elif CI.H is not None and CO.H is None and self.CapacityD is not None:
                H = CI.H + self.CapacityD / (mc * dt)
                self.Cold_Out = self.Model.Prop(
                    CO.fluid, StatePointName=CO.StatePointName,
                    P=CO.P, H=H, Mass_flowrate=mc)

            elif CI.H is not None and CO.H is not None and self.CapacityD is None:
                self.CapacityD = (CO.H - CI.H) * mc * dt

            # [FIX-6] Over-specified: CI.H, CO.H, and CapacityD all known → check
            elif CI.H is not None and CO.H is not None and self.CapacityD is not None:
                Q_check = (CO.H - CI.H) * mc * dt
                tol = 0.01 * max(abs(self.CapacityD), 1.0)
                if abs(Q_check - self.CapacityD) > tol:
                    raise ValueError(
                        f"[{self.ID}] TES Discharging over-specified and inconsistent: "
                        f"enthalpy-based capacity = {Q_check:.2f} J, "
                        f"given CapacityD = {self.CapacityD:.2f} J."
                    )

            elif CI.H is None and CO.H is None and self.CapacityD is not None:
                self.Cold_Out = self.Model.Prop(
                    CO.fluid, StatePointName=CO.StatePointName,
                    P=CO.P, T=self.T_melt - self.PPT, Mass_flowrate=mc)
                H = self.Cold_Out.H - self.CapacityD / (mc * dt)
                self.Cold_In = self.Model.Prop(
                    CI.fluid, StatePointName=CI.StatePointName,
                    P=CI.P, H=H, Mass_flowrate=mc)

            elif CI.H is None and CO.H is not None and self.CapacityD is not None:
                H = CO.H - self.CapacityD / (mc * dt)
                self.Cold_In = self.Model.Prop(
                    CI.fluid, StatePointName=CI.StatePointName,
                    P=CI.P, H=H, Mass_flowrate=mc)

            else:
                raise ValueError(
                    f"[{self.ID}] TES Discharging: unresolvable combination of "
                    f"CI.H={CI.H}, CO.H={CO.H}, CapacityD={self.CapacityD}. "
                    f"Provide at least one enthalpy and capacity, or both enthalpies."
                )

            self.Discharging_Power = self.CapacityD / dt

        # --- Charging ---
        elif self.Charge == 'Charging':
            HI = self.Hot_In
            HO = self.Hot_Out
            mh = self.Hot_Mass_flowrate
            dt = self.Charging_time * 3600

            if HI.H is not None and HO.H is None and self.Capacity is None:
                if self.PPT < (HI.T - self.T_melt):
                    self.Hot_Out = self.Model.Prop(
                        HO.fluid, StatePointName=HO.StatePointName,
                        P=HO.P, T=self.T_melt + self.PPT, Mass_flowrate=mh)
                    self.Capacity = (HI.H - self.Hot_Out.H) * mh * dt
                else:
                    raise ValueError(
                        f"[{self.ID}] Charging: (T_melt + PPT) = "
                        f"{self.T_melt + self.PPT:.2f} K ≥ T_hot_in = {HI.T:.2f} K. "
                        f"Reduce PPT or lower T_melt."
                    )

            elif HI.H is not None and HO.H is None and self.Capacity is not None:
                H = HI.H - self.Capacity / (mh * dt)
                self.Hot_Out = self.Model.Prop(
                    HO.fluid, StatePointName=HO.StatePointName,
                    P=HO.P, H=H, Mass_flowrate=mh)

            elif HI.H is not None and HO.H is not None and self.Capacity is None:
                self.Capacity = (HI.H - HO.H) * mh * dt

            # [FIX-6] Over-specified: HI.H, HO.H, and Capacity all known → check
            elif HI.H is not None and HO.H is not None and self.Capacity is not None:
                Q_check = (HI.H - HO.H) * mh * dt
                tol = 0.01 * max(abs(self.Capacity), 1.0)
                if abs(Q_check - self.Capacity) > tol:
                    raise ValueError(
                        f"[{self.ID}] TES Charging over-specified and inconsistent: "
                        f"enthalpy-based capacity = {Q_check:.2f} J, "
                        f"given Capacity = {self.Capacity:.2f} J."
                    )

            elif HI.H is None and HO.H is None and self.Capacity is not None:
                self.Hot_Out = self.Model.Prop(
                    HO.fluid, StatePointName=HO.StatePointName,
                    P=HO.P, T=self.T_melt + self.PPT, Mass_flowrate=mh)
                H = self.Hot_Out.H + self.Capacity / (mh * dt)
                self.Hot_In = self.Model.Prop(
                    HI.fluid, StatePointName=HI.StatePointName,
                    P=HI.P, H=H, Mass_flowrate=mh)

            elif HI.H is None and HO.H is not None and self.Capacity is not None:
                H = HO.H + self.Capacity / (mh * dt)
                self.Hot_In = self.Model.Prop(
                    HI.fluid, StatePointName=HI.StatePointName,
                    P=HI.P, H=H, Mass_flowrate=mh)

            else:
                raise ValueError(
                    f"[{self.ID}] TES Charging: unresolvable combination of "
                    f"HI.H={HI.H}, HO.H={HO.H}, Capacity={self.Capacity}. "
                    f"Provide at least one enthalpy and capacity, or both enthalpies."
                )

            self.Charging_Power = self.Capacity / dt

        else:
            raise ValueError(
                f"Invalid Charge='{self.Charge}' in {self.ID}. "
                f"Use 'Charging' or 'Discharging'."
            )

        # exergy [FIX-10]
        try:
            self.Ex_D = ((self.Hot_In.Ex + self.Cold_In.Ex)
                         - (self.Hot_Out.Ex + self.Cold_Out.Ex))
        except Exception:
            self.Ex_D = "Not Calculated"

        self.Solution_Status = True

        if self.Hot_In_state:   self.Model.Point[self.Hot_In_state]   = self.Hot_In
        if self.Hot_Out_state:  self.Model.Point[self.Hot_Out_state]  = self.Hot_Out
        if self.Cold_In_state:  self.Model.Point[self.Cold_In_state]  = self.Cold_In
        if self.Cold_Out_state: self.Model.Point[self.Cold_Out_state] = self.Cold_Out

    def __str__(self):
        return (
            f"{self.ID} (TES - {self.Charge}):\n"
            f"  Charging Power   : {self.Charging_Power} W\n"
            f"  Discharging Power: {self.Discharging_Power} W\n"
            f"  Capacity         : {self.Capacity} J\n"
            f"  Solved           : {self.Solution_Status}"
        )
