import numpy as np
try:
    import pymc as pm
    from pytensor.tensor import where
except:
    pass

PRESSURE = 101325

# 计算含湿量
def get_HumRatio_from_Dew(dew_point,engine):
    """
    :return: 含湿量 kg 水蒸气 / kg 干空气
    """
    if engine == 'pymc':
        EXP = pm.math.exp
    elif engine == 'numpy':
        EXP = np.exp
    
    vapor_pressure = 6.112 * EXP(17.67 * dew_point / (dew_point + 243.5)) * 100
    HumRatio       = (0.622 * vapor_pressure) / (PRESSURE - vapor_pressure)
    return HumRatio

def get_Dew_from_HumRatio(HumRatio,engine):
    """
    :param HumRatio: 含湿量 kg 水蒸气 / kg 干空气
    :return: 露点温度
    """
    if engine == 'pymc':
        LOG = pm.math.log
        WHERE = where
    elif engine == 'numpy':
        LOG = np.log
        WHERE = np.where
    HumRatio       = WHERE(HumRatio<1e-5,1e-5,HumRatio)
    vapor_pressure = (HumRatio * PRESSURE) / (0.622 + HumRatio)
    x              = LOG(vapor_pressure / (6.112 * 100))
    dew_point      = (243.5 * x) / (17.67 - x)
    return dew_point

def get_Enthalpy_from_Tdb_and_HumRatio(Tdb,Hr,engine):
    # return: 焓值 (kJ/kg干空气)
    # 常数
    c_pa     = 1.006  # 干空气比热容 (kJ/kg·K)
    c_pv     = 1.805  # 水蒸气比热容 (kJ/kg·K)
    h_fg     = 2501   # 水的汽化潜热 (kJ/kg)
    Enthalpy = c_pa * Tdb + Hr * (h_fg + c_pv * Tdb)
    return Enthalpy

def get_Enthalpy_from_Tdb_and_Dew(Tdb,Dew,engine):
    Hr       = get_HumRatio_from_Dew(Dew,engine)
    Enthalpy = get_Enthalpy_from_Tdb_and_HumRatio(Tdb,Hr,engine)
    return Enthalpy

def get_mixed_Dew(F1,F2,Dew1,Dew2,engine):
    
    if engine == 'pymc':
        LOG = pm.math.log
        EXP = pm.math.exp
    elif engine == 'numpy':
        LOG = np.log
        EXP = np.exp
    
    # Antoine 方程计算饱和水蒸气压力
    def saturation_pressure(T):
        return 611 * EXP(17.27 * T / (T + 237.3))
    # 计算湿度比
    def humidity_ratio(e_s):
        return 0.622 * e_s / (PRESSURE - e_s)
    # 计算混合后的露点温度
    def dew_point_temperature(e_mix):
        return (237.3 * LOG(e_mix / 611)) / (17.27 - LOG(e_mix / 611))

    # 计算空气 A 和空气 B 的饱和水蒸气压力
    e_s1 = saturation_pressure(Dew1)
    e_s2 = saturation_pressure(Dew2)

    # 计算空气 A 和空气 B 的湿度比
    W1 = humidity_ratio(e_s1)
    W2 = humidity_ratio(e_s2)

    W_mix = (F1 * W1 + F2 * W2) / (F1 + F2) # 计算混合后的湿度比
    e_mix = (W_mix * PRESSURE) / (0.622 + W_mix)   # 计算混合后的水蒸气分压
    d_mix = dew_point_temperature(e_mix)    # 计算混合后的露点温度

    return d_mix

def get_RH_from_Tdb_and_Hr(Tdb, Hr, engine):
    """
    计算相对湿度 (0~1)，基于 ASHRAE 标准的高精度方法
    
    参数:
        Tdb: 干球温度 (°C)
        Hr: 绝对湿度 (kg/kg)
        P: 大气压力 (Pa)，默认 101325 Pa (1 atm)
    
    返回:
        RH: 相对湿度 (0~1)
    """
    if engine == 'pymc':
        LOG = pm.math.log
        EXP = pm.math.exp
    elif engine == 'numpy':
        LOG = np.log
        EXP = np.exp
    # 1. 计算饱和水蒸气压力 (Pa) - Hyland-Wexler 公式 (ASHRAE 标准)
    T_kelvin = Tdb + 273.15  # 转换为开尔文温度
    # 饱和水蒸气压力 (Pa) - Hyland-Wexler (1983)
    ln_es = (
        -5.8002206e3 / T_kelvin
        + 1.3914993
        - 4.8640239e-2 * T_kelvin
        + 4.1764768e-5 * T_kelvin**2
        - 1.4452093e-8 * T_kelvin**3
        + 6.5459673 * LOG(T_kelvin)
    )
    es = EXP(ln_es)
    # 2. 计算实际水蒸气压力 (Pa)
    e = Hr * PRESSURE / (0.621945 + Hr)  # 0.621945 ≈ 分子量比 (18.01528 / 28.966)
    # 3. 计算相对湿度 (RH)
    RH = e / es
    return RH