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

from ._base_components import BaseComponents
from ..tools.enthalpy import (
    get_Dew_from_HumRatio,
    get_HumRatio_from_Tdb_and_RH,
    get_RH_from_Tdb_and_Hr,
)


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_P6 = param['beta_P6']
        beta_P7 = param['beta_P7']
        beta_P8 = param['beta_P8']
        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 + beta_P6 * TinC + beta_P7 * TinP) 
            * EXP(-beta_P5 * FP) 
            + beta_P8 * FR
            + 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
        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    = 0
        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.Normal(f'{self.name}_beta_P2',sigma=1,initval=0),
            'beta_P3': pm.TruncatedNormal(f'{self.name}_beta_P3',mu=1,sigma=2,initval=1.5,lower=0),
            'beta_P4': pm.HalfNormal(f'{self.name}_beta_P4',sigma=1,initval=0.1),
            'beta_P5': pm.TruncatedNormal(f'{self.name}_beta_P5',mu=5,sigma=2,initval=5,lower=0),
            'beta_P6': pm.Normal(f'{self.name}_beta_P6',mu=0,sigma=1,initval=0),
            'beta_P7': pm.Normal(f'{self.name}_beta_P7',mu=0,sigma=1,initval=0),
            'beta_P8': pm.Normal(f'{self.name}_beta_P8',mu=1,sigma=10,initval=1),
            '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

class WheelS3V2(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_P6 = param['beta_P6']
        beta_P7 = param['beta_P7']
        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 * (1 - EXP(-beta_P5 * FR)) 
                + beta_P6 * TinC * (1 - EXP(-beta_P5 * FC))
                + beta_P7 * TinP * (1 - 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
        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    = 0
        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.Normal(f'{self.name}_beta_P2',sigma=1,initval=0),
            'beta_P3': pm.TruncatedNormal(f'{self.name}_beta_P3',mu=1,sigma=2,initval=1.5,lower=0),
            'beta_P4': pm.HalfNormal(f'{self.name}_beta_P4',sigma=1,initval=0.1),
            'beta_P5': pm.HalfNormal(f'{self.name}_beta_P5',sigma=10,initval=1),
            'beta_P6': pm.Normal(f'{self.name}_beta_P6',mu=0,sigma=1,initval=0),
            'beta_P7': pm.Normal(f'{self.name}_beta_P7',mu=0,sigma=1,initval=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

class WheelS3V3(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_P6 = param['beta_P6']
        beta_P7 = param['beta_P7']
        beta_C1 = param['beta_C1']
        beta_C2 = param['beta_C2']
        beta_C3 = param['beta_C3']
        
        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 * (1 - EXP(-beta_P5 * FR)) 
                + beta_P6 * TinC * (1 - EXP(-beta_P5 * FC))
                + beta_P7 * TinP * (1 - 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
        TdiffP  = (Q_lat_P + Q_sen_P) / (FP * cls.CONSTANT['c_p_air'])
        ToutP   = TinP + TdiffP
        
        # 冷却侧
        TdiffC    = beta_C1 * RinC ** beta_C2 * TinR + beta_C3
        ToutC     = TinC + TdiffC
        HdiffC    = 0
        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.Normal(f'{self.name}_beta_P2',sigma=1,initval=0),
            'beta_P3': pm.TruncatedNormal(f'{self.name}_beta_P3',mu=1,sigma=2,initval=1.5,lower=0),
            'beta_P4': pm.HalfNormal(f'{self.name}_beta_P4',sigma=1,initval=0.1),
            'beta_P5': pm.HalfNormal(f'{self.name}_beta_P5',sigma=10,initval=1),
            'beta_P6': pm.Normal(f'{self.name}_beta_P6',mu=0,sigma=1,initval=0),
            'beta_P7': pm.Normal(f'{self.name}_beta_P7',mu=0,sigma=1,initval=0),
            'beta_C1': pm.TruncatedNormal(f'{self.name}_beta_C1',mu=0.5,sigma=1,initval=0.5,lower=0),
            'beta_C2': pm.TruncatedNormal(f'{self.name}_beta_C2',mu=0.3,sigma=1,initval=0.3,lower=0),
            'beta_C3': pm.Normal(f'{self.name}_beta_C3',mu=0,sigma=10,initval=0),
        }
        return param