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_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_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    = (beta_P1 * RinC**beta_P4 * HinC * TinR * EXP(-beta_P5 * FC) + beta_P2)/1000
        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_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']
        
        K0 = param['K0']
        eta = param['eta']
        beta_P1 = param['beta_P1']
        beta_P2 = param['beta_P2']
        beta_P3 = param['beta_P3']
        beta_P1a = param['beta_P1a']
        beta_P2a = param['beta_P2a']
        beta_P3a = param['beta_P3a']
        beta_P4 = param['beta_P4']
        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)
        T_wheel_P = (
            beta_P1a * TinP * FP ** beta_P1 
            + beta_P2a * TinC * FC ** beta_P2
            + beta_P3a * TinR * FR ** beta_P3
        ) 
        E = -45
        R = 8.314
        N = 0.45
        H_eq    = K0 * EXP(-E / (R * (T_wheel_P + 273.15))) * RinP ** N # 处理侧平衡状态下转轮能够吸附的水
        HoutP   = HinP - eta * (H_eq - HinP) * (1 - EXP(-beta_P4 * FP))
        DoutP   = get_Dew_from_HumRatio(HoutP,engine)
        WdiffP  = (HinP - HoutP) * FP
        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.Normal(f'{self.name}_beta_P1',mu=0,sigma=1,initval=0),
            'beta_P2': pm.Normal(f'{self.name}_beta_P2',mu=0,sigma=1,initval=0),
            'beta_P3': pm.Normal(f'{self.name}_beta_P3',mu=0,sigma=1,initval=0),
            'beta_P1a': pm.Normal(f'{self.name}_beta_P1a',mu=1,sigma=1,initval=1),
            'beta_P2a': pm.Normal(f'{self.name}_beta_P2a',mu=1,sigma=1,initval=1),
            'beta_P3a': pm.Normal(f'{self.name}_beta_P3a',mu=1,sigma=1,initval=1),
            'beta_P4': pm.TruncatedNormal(f'{self.name}_beta_P4',mu=5,sigma=3,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),
            'K0'     : pm.TruncatedNormal(f'{self.name}_k0',mu=0.4,sigma=0.1,initval=0.4,lower=0),
            'eta'    : pm.HalfNormal(f'{self.name}_eta',sigma=0.1,initval=0.1)
        }
        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'] # km
        beta_P2 = param['beta_P2'] # eta_max
        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']
        beta_R1 = param['beta_R1']
        beta_R2 = param['beta_R2']
        
        # FP_w = EXP(FP)
        # FR_w = EXP(FR)
        # FC_w = EXP(FC)
        T_wheel = (
            TinP * (FP / (FP + FR + FC))
            + TinR * (FR / (FP + FR + FC))
            + TinC * (FC / (FP + FR + FC))
        )
        Hr_eq  = get_HumRatio_from_Tdb_and_RH(T_wheel,1,engine)
        
        # 处理侧
        # 湿度
        eta    = beta_P2 * (1-EXP(-beta_P1 * (HinP-Hr_eq)/Hr_eq))
        HdiffP = eta * (HinP - Hr_eq)
        HoutP  = HinP - HdiffP
        DoutP  = get_Dew_from_HumRatio(HoutP,engine)
        # 温度
        TdiffP_sen = beta_P4 * (T_wheel - TinP)
        TdiffP_lat = cls.CONSTANT['h_ads'] / cls.CONSTANT['c_p_air'] * HdiffP * beta_P5
        ToutP      = TinP + TdiffP_sen + TdiffP_lat
        
        # 预冷侧
        # 湿度
        HdiffC = eta * (HinC - Hr_eq)
        HoutC  = HinC - HdiffC
        DoutC  = get_Dew_from_HumRatio(HoutC,engine)
        # 温度
        TdiffC = beta_C1 * EXP(-beta_C2 * EXP(-beta_C3 * (T_wheel - TinC))) * EXP(-beta_C4 * FC)
        ToutC  = TinC + TdiffC
        
        # 再生侧
        # 湿度
        HdiffR = (HdiffP * FP + HdiffC * FC) / FR
        HoutR  = HinR + HdiffR
        DoutR  = get_Dew_from_HumRatio(HoutR,engine)
        # 温度
        TdiffR_lat = cls.CONSTANT['h_ads'] / cls.CONSTANT['c_p_air'] * HdiffR * beta_R1
        TdiffR_sen = beta_R2 * (TinR - T_wheel)
        ToutR      = TinR - TdiffR_sen - TdiffR_lat
        
        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.HalfNormal(f'{self.name}_beta_P1',sigma=10,initval=0.01),
            'beta_P2': pm.Uniform(f'{self.name}_beta_P2',lower=0,upper=1,initval=0.5),
            'beta_P4': pm.HalfNormal(f'{self.name}_beta_P4', sigma=0.5,initval=0.2),
            'beta_P5': pm.TruncatedNormal(f'{self.name}_beta_P5', mu=1, sigma=1,initval=1,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),
            'beta_R1': pm.TruncatedNormal(f'{self.name}_beta_R1', mu=1, sigma=1,initval=1,lower=0),
            'beta_R2': pm.HalfNormal(f'{self.name}_beta_R2', sigma=10,initval=1),
        }
        return param