import numpy as np
try:
    import pymc as pm
except:
    pass

from ._base_components import BaseComponents
from ..tools.enthalpy import get_Dew_from_HumRatio
from ..tools.enthalpy import get_HumRatio_from_Dew


class CoolingCoil(BaseComponents):
    
    def __init__(self, name):
        super().__init__(name)
    
    @classmethod
    def model(
        cls,
        TinA,HinA,FA,
        TinW,FW,
        engine,
        param
    ):
        FUNC  = cls.get_func_by_engine(engine)
        WHERE = FUNC['WHERE']
        GT    = FUNC['GT']
        EXP   = FUNC['EXP']
        
        DinA = get_Dew_from_HumRatio(HinA,engine)
        mode = WHERE(GT(DinA-TinW,0),1,0) #1湿工况
        
        UA_dry       = param['UA_dry']
        UA_wet_ratio = param['UA_wet_ratio']
        Ts_adj       = param['Ts_adj']
        eta          = param['eta']
        UA_wet       = UA_dry * UA_wet_ratio
        
        C_min   = FA * cls.CONSTANT['c_p_air']
        C_max   = FW * cls.CONSTANT['c_p_water']
        Cr      = C_min / C_max
        NTU_wet = UA_wet / C_min
        NTU_dry = UA_dry / C_min
        NTU     = WHERE(mode,NTU_wet,NTU_dry)
        epsilon = (1-EXP(-NTU*(1-Cr))) / (1-Cr*EXP(-NTU*(1-Cr)))
        
        # 出风干球温度
        Ts    = TinW + Ts_adj * (HinA) * (TinA - TinW) # 湿工况盘管表面温度
        Q_wet = epsilon * C_min * (TinA - Ts)
        Q_dry = epsilon * C_min * (TinA - TinW)
        Q     = WHERE(mode,Q_wet,Q_dry)
        ToutA = TinA - Q / (FA * cls.CONSTANT['c_p_air'])
        
        # 出风湿度+露点
        HoutA = WHERE(
            mode,
            HinA - eta * (HinA - get_HumRatio_from_Dew(TinW,engine)),
            HinA
        )
        DoutA = get_Dew_from_HumRatio(HoutA,engine)
        
        return {'ToutA':ToutA,'HoutA':HoutA,'DoutA':DoutA} 

    def prior(self) -> dict:
        param = {
            'UA_dry'      : pm.HalfNormal(f'{self.name}_UA_dry',5,initval=10),
            'UA_wet_ratio': pm.TruncatedNormal(f'{self.name}_UA_wet_ratio',mu=4,sigma=0.2,lower=1,upper=10,initval=4),
            'Ts_adj'      : pm.HalfNormal(f'{self.name}_Ts_adj',5,initval=5),
            'eta'         : pm.HalfNormal(f'{self.name}_eta',0.5,initval=0.5),
        }
        return param


class CoolingCoil2(BaseComponents):
    
    def __init__(self, name):
        super().__init__(name)
    
    @classmethod
    def model(
        cls,
        TinA,HinA,FA,
        TinW,FW,
        engine,
        param
    ):
        FUNC  = cls.get_func_by_engine(engine)
        WHERE = FUNC['WHERE']
        EXP   = FUNC['EXP']
        
        DinA = get_Dew_from_HumRatio(HinA,engine)
        
        beta1_T     = param['beta1_T']
        beta2_T     = param['beta2_T']
        beta3_T     = param['beta3_T']
        beta1_D_adj = param['beta1_D_adj']
        beta2_D_adj = param['beta2_D_adj']
        beta3_D_adj = param['beta3_D_adj']
        beta4_D     = param['beta4_D']
        
        beta1_D = beta1_T * beta1_D_adj
        beta2_D = beta2_T * beta2_D_adj
        beta3_D = beta3_T * beta3_D_adj
        
        EFF_T_FW = 1 - EXP(-beta2_T * FW)
        EFF_T_FA = EXP(-beta3_T * FA)
        EFF_D_FW = 1 - EXP(-beta2_D * FW)
        EFF_D_FA = EXP(-beta3_D * FA)
        mode     = WHERE(DinA > TinW,1,0)
        ToutA    = TinA - beta1_T * (TinA - TinW) * EFF_T_FA * EFF_T_FW
        DoutA    = DinA - (beta1_D - beta4_D * (TinA - DinA)) * (DinA - TinW) * mode * EFF_D_FW * EFF_D_FA
        HoutA    = get_HumRatio_from_Dew(DoutA,engine)
        
        return {'ToutA':ToutA,'HoutA':HoutA,'DoutA':DoutA,'FA':FA} 

    def prior(self) -> dict:
        param = {
            'beta1_T'    : pm.TruncatedNormal(f'{self.name}_beta1_T',mu=1,sigma=0.2,lower=0),
            'beta2_T'    : pm.TruncatedNormal(f'{self.name}_beta2_T',mu=5,sigma=3,lower=0),
            'beta3_T'    : pm.TruncatedNormal(f'{self.name}_beta3_T',mu=5,sigma=0.3,lower=0),
            'beta1_D_adj': pm.TruncatedNormal(f'{self.name}_beta1_D_adj',mu=1,sigma=0.1,lower=0.1),
            'beta2_D_adj': pm.TruncatedNormal(f'{self.name}_beta2_D_adj',mu=1,sigma=0.1,lower=0.1),
            'beta3_D_adj': pm.TruncatedNormal(f'{self.name}_beta3_D_adj',mu=1,sigma=0.1,lower=0.1),
            'beta4_D'    : pm.TruncatedNormal(f'{self.name}_beta4_D',mu=0.01,sigma=0.01,lower=0)
        }
        return param