# -*- coding: utf-8 -*-
"""Component models used by Brayton cycle simulations."""


class ComponentMassMixin:
    """Shared mass-result storage interface for cycle components."""

    component_type = "component"

    def _init_mass_interface(self):
        self.mass = None
        self.mass_variables = None

    def set_mass_result(self, result):
        """Store a mass-model result on this component.

        Mass formulas live in ``brayton_cycle.mass_models``. Components only
        keep the result after their thermodynamic calculator has populated
        ``variables``.
        """
        if self.variables is None:
            raise ValueError("Run component calculator before storing mass")

        if isinstance(result, dict):
            if "mass" in result:
                mass = result["mass"]
            elif "total_mass" in result:
                mass = result["total_mass"]
            else:
                raise ValueError("Mass result dictionary must include 'mass'")
            mass_variables = dict(result)
        else:
            mass = result
            mass_variables = {"mass": result}

        self.mass = mass
        self.mass_variables = mass_variables
        self.variables["mass"] = self.mass
        self.variables["mass_variables"] = self.mass_variables
        return self.mass

    def clear_mass_result(self):
        self.mass = None
        self.mass_variables = None
        if self.variables is not None:
            self.variables.pop("mass", None)
            self.variables.pop("mass_variables", None)

class Compressor(ComponentMassMixin):
    """压缩机类"""
    component_type = "compressor"

    def __init__(self, name, eff):
        """
        初始化参数
        
        name: 名称
        eff: 等熵效率
        Wc: 压缩功
        inlet_state: 入口参数
        outlet_state: 出口参数
        outlet_state_is: 等熵状态下出口参数
        """
        
        self.name = name
        self.eff = eff
        self.variables = None
        self._init_mass_interface()
        
    def calculator(self, p_in, T_in, p_out, property_calculator):
        # 先计算熵值
        inlet_state = property_calculator.calculate_properties(T=T_in, P=p_in)
        mw = property_calculator.mw
        s = inlet_state['s']
        h_in = inlet_state['h']
        outlet_state_is = property_calculator.calculate_properties(P=p_out, s=s)
        h_out_is = outlet_state_is['h']
        h_out = h_in + (h_out_is - h_in) / self.eff
        outlet_state = property_calculator.calculate_properties(P=p_out, h=h_out)
        Wc = h_out - h_in
        T_out = outlet_state['T']
        # 计算结果
        self.variables = {
            'name': self.name,
            'inlet_state':{
                'P': p_in,          # 压强(kPa)
                'T': T_in,
                'h': h_in/mw,       # 比焓(J/mol)->(J/kg)
                's': s/mw,          # 比熵(J/mol.K)->(J/kg.K)
                },
            'outlet_state':{
                'P': p_out,         # 压强(kPa)
                'T': T_out,
                'h': h_out/mw,      # 比焓(kJ/mol)->(kJ/kg)
                's': s/mw,          # 比熵(kJ/mol.K)->(kJ/kg.K)
                },
            'eff': self.eff,
            'Wc': Wc/mw,       # 压缩功(kJ/mol)->(kJ/kg)
            'pi': p_out / p_in
            }

class Turbine(ComponentMassMixin):
    """透平类"""
    component_type = "turbine"
    def __init__(self, name, eff):
        """
        初始化参数
        name: 名称
        eff: 透平效率
        """
        self.name = name
        self.eff = eff
        self.variables = None
        self._init_mass_interface()
    
    def calculator(self, p_in, T_in, p_out, property_calculator):
        """涡轮参数计算"""
        inlet_state = property_calculator.calculate_properties(P=p_in, T=T_in)
        mw = property_calculator.mw
        s = inlet_state['s']
        h_in = inlet_state['h']
        outlet_state_is = property_calculator.calculate_properties(P=p_out, s=s)
        h_out_is = outlet_state_is['h']
        h_out = h_in - (h_in - h_out_is) * self.eff
        outlet_state = property_calculator.calculate_properties(P=p_out, h=h_out)
        T_out = outlet_state['T']
        Wt = h_in - h_out
        self.variables = {
            'name': self.name,
            'inlet_state':{
                'P': p_in,          # 压强(kPa)
                'T': T_in,
                'h': h_in/mw,      # 比焓(J/mol)->(J/kg)
                's': s/mw,         # 比熵(J/mol.K)->(J/kg.K)
                },
            'outlet_state':{
                'P': p_out,        # 压强(kPa)
                'T': T_out,
                'h': h_out/mw,     # 比焓(J/mol)->(J/kg)
                's': s/mw,         # 比熵(J/mol.K)->(J/kg.K)
                },
            'eff': self.eff,
            'Wt': Wt/mw,      # 透平做功(J/mol)->(J/kg)
            'pi': p_in / p_out
            }

class Recuperator(ComponentMassMixin):
    """换热器类"""
    component_type = "recuperator"
    def __init__(self, name, eff, x=0):
        """
        初始化参数
        name: 名称
        eff: 换热效率(基于焓的计算方法)
        """
        self.name = name
        self.eff = eff
        self.Q_ex = None
        self.variables = None
        self.x = x
        self._init_mass_interface()
        
    def calculator(self, cold_inlet_state, hot_inlet_state, bypass_info, ploss=0.0, property_calculator=None):
        """
        计算换热器两侧参数，默认逆流
        bypass_info: 是否存在分流，0为不存在，1为存在
        
        下标含义
        ----------
        1: 换热器冷端入口
        2: 换热器冷端出口
        3: 换热器热端入口
        4: 换热器热端出口
        ass: 迭代中间变量，假设值
        """
        
        # 计算两入口参数, 这里单位是kg
        if property_calculator is None:
            if hasattr(ploss, "calculate_properties"):
                property_calculator = ploss
                ploss = 0.0
            else:
                raise ValueError("property_calculator is required")

        p1 = cold_inlet_state['P']
        T1 = cold_inlet_state['T']
        h1 = cold_inlet_state['h']
        p2 = p1 * (1 - ploss)
        p3 = hot_inlet_state['P']
        T3 = hot_inlet_state['T']
        h3 = hot_inlet_state['h']
        p4 = p3 * (1 - ploss)
        mw = property_calculator.mw
        
        # 设定质量流量
        m_cold = (1-self.x) if bypass_info == 1 else 1.0
        m_hot = 1.0
        
# =============================================================================
#         # 采用焓差效能的方式来计算换热器进出口参数
#         # 假设最大温差发生在冷端, 计算冷端出口温度和比焓
#         cold_outlet_state = property_calculator.calculate_properties(T=T3, P=p2)
#         h2 = cold_outlet_state['h'] / mw
#         Q_ass_cold = m_cold * abs(h2 - h1)
#         
#         # 假设最大温差发生在热端, 计算热端出口温度和比焓
#         hot_outlet_state = property_calculator.calculate_properties(T=T1, P=p4)
#         h4 = hot_outlet_state['h'] / mw
#         Q_ass_hot = m_hot * abs(h3 - h4)
#         
#         # 比较两个可能的Q，取最小值与焓差效能的乘积作为实际换热量
#         self.Q_ex = min(Q_ass_hot, Q_ass_cold) * self.eff
#         # 由实际换热量计算出口焓和出口状态
#         h2 = h1 + self.Q_ex / m_cold
#         h4 = h3 - self.Q_ex / m_hot
#         cold_outlet_state = property_calculator.calculate_properties(P=p2, h=h2*mw)
#         hot_outlet_state = property_calculator.calculate_properties(P=p4, h=h4*mw)
#         T2 = cold_outlet_state['T']
#         T4 = hot_outlet_state['T']
# =============================================================================

        # 采用温差效能的方式来计算换热器进出口参数
        T_ass_max = abs(T1 - T3)
        # 假设最大温差发生在冷端, 计算冷端出口温度和比焓
        T2 = T1 + T_ass_max * self.eff
        cold_outlet_state = property_calculator.calculate_properties(T=T2, P=p2)
        h2 = cold_outlet_state['h'] / mw
        Q_ass_cold = m_cold * abs(h2 - h1)
        
        # 假设最大温差发生在热端, 计算热端出口温度和比焓
        T4 = T3 - T_ass_max * self.eff
        hot_outlet_state = property_calculator.calculate_properties(T=T4, P=p4)
        h4 = hot_outlet_state['h'] / mw
        Q_ass_hot = m_hot * abs(h3 - h4)
        
        # 比较两个可能的Q，取最小值与焓差效能的乘积作为实际换热量
        self.Q_ex = min(Q_ass_hot, Q_ass_cold)
        # 由实际换热量计算出口焓和出口状态
        h2 = h1 + self.Q_ex / m_cold
        h4 = h3 - self.Q_ex / m_hot
        cold_outlet_state = property_calculator.calculate_properties(P=p2, h=h2*mw)
        hot_outlet_state = property_calculator.calculate_properties(P=p4, h=h4*mw)
        T2 = cold_outlet_state['T']
        T4 = hot_outlet_state['T']
        
        # 拼装变量
        self.variables = {
            'name': self.name,
            'cold_inlet_state':{
                'P': p1,
                'T': T1,
                'h': h1,
                's': cold_inlet_state['s']
                },
            'cold_outlet_state':{
                'P': p2,
                'T': cold_outlet_state['T'],
                'h': h2,
                's': cold_outlet_state['s']/mw 
                },
            'hot_inlet_state':{
                'P': p3,
                'T': T3,
                'h': h3,
                's': hot_inlet_state['s']
                },
            'hot_outlet_state':{
                'P': p4,
                'T': T4,
                'h': h4,
                's': hot_outlet_state['s']/mw
                },
            'eff': self.eff,
            'Q_exchange': self.Q_ex
            }
    def check_pinch_point(self, property_calculator, num_segments=20):
        """
        换热器内部夹点校验
        将换热量均分为 num_segments 段，检查内部每个微元的冷热流体温度
        """
        h_cold_in = self.variables['cold_inlet_state']['h']
        h_hot_in = self.variables['hot_inlet_state']['h']
        p_cold = self.variables['cold_inlet_state']['P']
        p_hot = self.variables['hot_inlet_state']['P']
        
        m_cold = (1 - self.x) if self.name == "Low Temperature recuprerator" else 1.0
        m_hot = 1.0
        
        dQ = self.Q_ex / num_segments
        
        # 沿冷流体流动方向步进检查
        for i in range(num_segments + 1):
            q_current = i * dQ
            
            # 当前微元截面的焓值
            h_cold_local = h_cold_in + q_current / m_cold
            h_hot_local = (h_hot_in - self.Q_ex / m_hot) + q_current / m_hot
            
            # 查温度
            T_cold_local = property_calculator.calculate_properties(P=p_cold, h=h_cold_local * property_calculator.mw)['T']
            T_hot_local = property_calculator.calculate_properties(P=p_hot, h=h_hot_local * property_calculator.mw)['T']
            
            # 如果热流体温度低于等于冷流体温度 (设定一个 0.1K 的最小逼近温差容差)
            if T_hot_local - T_cold_local < 0.1:
                return False  # 发生温度交叉，物理不可行！
                
        return True

class Heater(ComponentMassMixin):
    """加热器类"""
    component_type = "heater"
    def __init__(self, name):
        self.name = name
        self.variables = None
        self._init_mass_interface()
    def calculator(self, inlet_state, outlet_state):
        h_in = inlet_state['h']
        h_out = outlet_state['h']
        Q_input = h_out - h_in
        self.variables = {
            'name': self.name,
            'inlet_state':{
                'P': inlet_state['P'],  # 压强(kPa)
                'T': inlet_state['T'],
                'h': inlet_state['h'],
                's': inlet_state['s'],
                },
            'outlet_state':{
                'P': outlet_state['P'], # 压强(kPa)
                'T': outlet_state['T'],
                'h': outlet_state['h'],
                's': outlet_state['s'],
                },
            'Q_in': Q_input,
            }

class Condenser(ComponentMassMixin):
    """冷凝器类"""
    component_type = "condenser"
    def __init__(self, name):
        self.name = name
        self.variables = None
        self._init_mass_interface()
    def calculator(self, inlet_state, outlet_state):
        h_in = inlet_state['h']
        h_out = outlet_state['h']
        Q_output = h_in - h_out
        self.variables = {
            'name': self.name,
            'inlet_state':{
                'P': inlet_state['P'],  # 压强(kPa)
                'T': inlet_state['T'],
                'h': inlet_state['h'],
                's': inlet_state['s'],
                },
            'outlet_state':{
                'P': outlet_state['P'], # 压强(kPa)
                'T': outlet_state['T'],
                'h': outlet_state['h'],
                's': outlet_state['s'],
                },
            'Q_out': Q_output,
            }

class Concentrator(ComponentMassMixin):
    """汇流组件"""
    component_type = "concentrator"
    def __init__(self, name):
        self.name = name
        self.variables = None
        self._init_mass_interface()
    def calculator(self, inlet_state_bypass, inlet_state_mroad, x, property_calculator):
        h_in_bypass = inlet_state_bypass['h']
        h_in_mroad = inlet_state_mroad['h']
        p_in = inlet_state_bypass['P']
        mw = property_calculator.mw
        h_out = x * h_in_bypass + (1-x) * h_in_mroad
        outlet_state = property_calculator.calculate_properties(P=p_in, h=h_out*mw)
        self.variables = {
            'name': self.name,
            'inlet_state_bypass': inlet_state_bypass,
            'inlet_state_mroad': inlet_state_mroad,
            'outlet_state':{
                'P': p_in,
                'T': outlet_state['T'],
                'h': h_out,
                's': outlet_state['s']/mw
                }
            }