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

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


class WheelS2(BaseComponents):
    
    def __init__(self, name):
        super().__init__(name)
    
    @classmethod
    def model(
        cls,
        TinP,HinP,FP,
        TinR,HinR,FR,
        engine  : str,
        param   : dict,
    ):
        FUNC = cls.get_func_by_engine(engine)
        EXP  = FUNC['EXP']
        RinP = get_RH_from_Tdb_and_Hr(TinP,HinP,engine)
        
        beta_Q1 = param['beta_Q1']
        beta_H1 = param['beta_H1']
        beta_H2 = param['beta_H2']
        beta_H3 = param['beta_H3']
        beta_H4 = param['beta_H4']
        
        # 转轮的温度
        T_avg = TinR * (1 - beta_H4) + TinP * beta_H4
        
        # 出风露点（除湿量）
        ## 处理侧
        Hdiff_mu = beta_H1 * RinP ** beta_H2 * EXP(-beta_H3 / T_avg) / 1000 # kg水蒸气/kg干空气
        HoutP_mu = HinP - Hdiff_mu
        DoutP_mu = get_Dew_from_HumRatio(HoutP_mu,engine)
        ## 再生侧
        HoutR_mu = HinR + Hdiff_mu * (FP / FR)
        DoutR_mu = get_Dew_from_HumRatio(HoutR_mu,engine)
        
        # 出风温度
        # 处理侧
        Q_latent   = Hdiff_mu * cls.CONSTANT['h_ads'] # Kj/kg 潜热
        Q_sensible = beta_Q1 * (TinR - TinP)          # Kj/kg 显热（影响升焓的部分）
        Q          = (Q_latent + Q_sensible) * FP     # Kj
        ToutP_mu   = TinP + Q / (FP * cls.CONSTANT['c_p_air'])
        # 再生侧
        ToutR_mu   = TinR - Q / (FR * cls.CONSTANT['c_p_air'])
        
        return {
            'ToutP':ToutP_mu,'HoutP':HoutP_mu,'DoutP':DoutP_mu,
            'ToutR':ToutR_mu,'HoutR':HoutR_mu,'DoutR':DoutR_mu,
        }
    
    def prior(self):
        param = {
            'beta_Q1': pm.HalfNormal(f'{self.name}_beta_Q1',sigma=10),
            'beta_H1': pm.HalfNormal(f'{self.name}_beta_H1',sigma=10),
            'beta_H2': pm.HalfNormal(f'{self.name}_beta_H2',sigma=10),
            'beta_H3': pm.HalfNormal(f'{self.name}_beta_H3',sigma=10),
            'beta_H4': pm.Uniform(f'{self.name}_beta_H4',lower=0,upper=1),
        }
        return param


class WheelS3(BaseComponents):
    def __init__(self, name):
        super().__init__(name)
    
    @classmethod
    def model(
        cls,
        TinP,HinP,FP,
        TinR,HinR,FR,
        TinC,HinC,FC,
        engine  : str,
        param   : dict,
    ):
        FUNC = cls.get_func_by_engine(engine)
        EXP  = FUNC['EXP']
        
        beta_P1 = param['beta_P1']
        beta_P2 = param['beta_P2']
        beta_P3 = param['beta_P3']
        beta_P4 = param['beta_P4']
        beta_P5 = param['beta_P5']
        beta_C1 = param['beta_C1']
        beta_C2 = param['beta_C2']
        beta_C3 = param['beta_C3']
        beta_C4 = param['beta_C4']
        
        RinP = get_RH_from_Tdb_and_Hr(TinP,HinP,engine)
        RinC = get_RH_from_Tdb_and_Hr(TinC,HinC,engine)
        
        # 处理侧
        HdiffP  = (beta_P1 * RinP**beta_P4 * HinP * TinR * EXP(-beta_P5 * FP) + beta_P2)/1000
        WdiffP  = HdiffP * FP
        HoutP   = HinP - HdiffP
        DoutP   = get_Dew_from_HumRatio(HoutP,engine)
        Q_lat_P = WdiffP * cls.CONSTANT['h_ads']
        Q_sen_P = beta_P3 * (TinR - TinP) * FP #TODO
        TdiffP  = (Q_lat_P + Q_sen_P) / (FP * cls.CONSTANT['c_p_air'])
        ToutP   = TinP + TdiffP
        
        # 冷却侧
        TdiffC    = beta_C1 * EXP(-beta_C2 * EXP(-beta_C3 * (TinR - TinC))) * EXP(-beta_C4 * FC)
        ToutC     = TinC + TdiffC
        HdiffC    = (beta_P1 * RinC**beta_P4 * HinC * TinR * EXP(-beta_P5 * FC) + beta_P2)/1000
        WdiffC    = HdiffC * FC
        HoutC     = HinC - HdiffC
        DoutC     = get_Dew_from_HumRatio(HoutC,engine)
        Q_total_C = TdiffC * FC * cls.CONSTANT['c_p_air']
        
        # 再生侧
        WdiffR    = WdiffP + WdiffC
        HoutR     = (HinR * FR + WdiffR) / FR
        DoutR     = get_Dew_from_HumRatio(HoutR,engine)
        Q_total_R = Q_lat_P + Q_sen_P + Q_total_C
        TdiffR    = Q_total_R / (FR * cls.CONSTANT['c_p_air'])
        ToutR     = TinR - TdiffR
        
        return {
            'ToutP':ToutP,'HoutP':HoutP,'DoutP':DoutP,'FP':FP,
            'ToutR':ToutR,'HoutR':HoutR,'DoutR':DoutR,'FR':FR,
            'ToutC':ToutC,'HoutC':HoutC,'DoutC':DoutC,'FC':FC,
        } 
        
    def prior(self):
        param = {
            'beta_P1': pm.TruncatedNormal(f'{self.name}_beta_P1',mu=5,sigma=10,initval=5,lower=0),
            'beta_P2': pm.TruncatedNormal(f'{self.name}_beta_P2',mu=0.5,sigma=1,initval=0.02,lower=0),
            'beta_P3': pm.TruncatedNormal(f'{self.name}_beta_P3',mu=1,sigma=2,initval=1.5,lower=0),
            'beta_P4': pm.TruncatedNormal(f'{self.name}_beta_P4',mu=1,sigma=0.3,initval=1,lower=0),
            'beta_P5': pm.TruncatedNormal(f'{self.name}_beta_P5',mu=5,sigma=2,initval=5,lower=0),
            'beta_C1': pm.TruncatedNormal(f'{self.name}_beta_C1',mu=60,sigma=10,initval=60,lower=10),
            'beta_C2': pm.TruncatedNormal(f'{self.name}_beta_C2',mu=30,sigma=10,initval=30,lower=1),
            'beta_C3': pm.TruncatedNormal(f'{self.name}_beta_C3',mu=0.05,sigma=0.1,initval=0.05,lower=0),
            'beta_C4': pm.TruncatedNormal(f'{self.name}_beta_C4',mu=1,sigma=1,initval=1,lower=0),
        }
        return param