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

from ._base_components import BaseComponents


class SteamCoilFs(BaseComponents):
    def __init__(self, name):
        super().__init__(name)
    
    @classmethod
    def model(
        cls,
        TinA,ToutA,FA,
        param,
        engine
    ):
        b1 = param['b1']
        b2 = param['b2']
        b3 = param['b3']
        Fs = b1 * (ToutA - TinA) * FA + b2 * TinA + b3
        return {'Fs':Fs}
    
    def prior(self):
        param = {
            'b1': pm.HalfNormal(f'{self.name}_b1',sigma=10),
            'b2': pm.Normal(f'{self.name}_b2',sigma=10),
            'b3': pm.Normal(f'{self.name}_b3',sigma=10),
        }
        return param

class SteamCoilFs2(BaseComponents):
    def __init__(self, name):
        super().__init__(name)
    
    @classmethod
    def model(cls,TinA,ToutA,FA,param,engine):
        b1 = param['b1']
        b2 = param['b2']
        FP = np.exp(b1 * (ToutA - TinA)*FA + b2)
        
        FUNC  = cls.get_func_by_engine(engine)
        WHERE = FUNC['WHERE']
        LT    = FUNC['LT']
        
        param_change_point    = param['change_point']
        param_b1              = param['a1']
        param_b2              = param['a2']
        param_b3              = param['a3']
        param_Ps_change_point = param_b1 + param_b2 * param_change_point
        Fs = WHERE(
            LT(FP,param_Ps_change_point * param_change_point),
            get_root(param_b1,param_b2,FP),
            get_root(param_b1+(param_b2-param_b3)*param_change_point,param_b3,FP)
        )
        return {'FP':FP,'Fs':Fs}

    def prior(self):
        param = {
            'b1'          : pm.Normal(f'{self.name}_b1',mu=0.08,sigma=0.1,initval=0.08),
            'b2'          : pm.Normal(f'{self.name}_b2',mu=3.65,sigma=0.02,initval=3.65),
            'change_point': pm.Normal(f'{self.name}_change_point',mu=104,sigma=1,initval=104),
            'a1'          : pm.Normal(f'{self.name}_a1',mu=17.5,sigma=1,initval=17.5),
            'a2'          : pm.Normal(f'{self.name}_a2',mu=0.2,sigma=0.1,initval=0.2),
            'a3'          : pm.Normal(f'{self.name}_a3',mu=1.5,sigma=0.1,initval=1.5),
        }
        return param

class SteamCoil(BaseComponents):
    def __init__(self, name):
        super().__init__(name)
    
    @classmethod
    def model(
        cls,
        TinA,ToutA,FA,
        param,
        engine
    ):
        Q = (ToutA - TinA) * FA * cls.CONSTANT['c_p_air']
        return {'Q':Q,'Fs':Q}
    
    def prior(self):
        param = {}
        return param


def get_root(b0,b1,FP):
    return (np.sqrt(b0**2+4*b1*FP)-b0)/(2*b1)


class SteamCoilVal(BaseComponents):
    def __init__(self, name,Fs_rated):
        super().__init__(name)
        self.Fs_rated = Fs_rated
        if self.Fs_rated is None:
            raise Exception('Fs_rated must be provided')
        
    def model(
        self,
        TinA,ToutA,FA,
        param,
        engine
    ):
        FUNC = self.get_func_by_engine(engine)
        EXP  = FUNC['EXP']
        Val = smooth_exp_transition(
            x          = (ToutA - TinA) / 100,
            s          = param['s'] - param['s_offset'] * TinA / 100,
            a          = param['a'],
            b          = param['b'],
            d          = param['d'],
            FUNC_where = FUNC['WHERE']
        ) * 100
        sigma = EXP((ToutA - TinA) / 100 * param['alpha2'] + param['alpha1'])
        Fs    = self.get_Fs_by_Val_and_Tin(Val,TinA,engine=engine)
        return {'Val':Val,'sigma':sigma,'Fs':Fs}
    
    def prior(self):
        param = {
            's'       : pm.Normal(f'{self.name}_s', mu=0, sigma=10,initval=0),
            'a'       : pm.HalfNormal(f'{self.name}_a', sigma=10,initval=1),
            'b'       : pm.Normal(f'{self.name}_b', mu=0, sigma=10,initval=0),
            'd'       : pm.HalfNormal(f'{self.name}_d', sigma=10,initval=1),
            's_offset': pm.HalfNormal(f'{self.name}_s_offset', sigma=1,initval=0.1),
            'alpha1'  : pm.Normal(f'{self.name}_alpha1', mu=0, sigma=10,initval=-5),
            'alpha2'  : pm.HalfNormal(f'{self.name}_alpha2', sigma=10,initval=5)
        }
        return param

    def get_Fs_by_Val_and_Tin(self,Val,Tin,engine):
        WHERE = self.get_func_by_engine(engine)['WHERE']
        
        def piecewise_function(x, a, b, s, k):
            y_max = a * s + b + a / k
            y = WHERE(
                x < s, 
                a * x + b, 
                y_max - (a / k) * np.exp(-k * (x - s))
            )
            return y
        
        Fs = piecewise_function(
            Val / 100,
            a = 4.48157826,
            b = -0.62511713,
            s = 0.47809967 + Tin / 100 * (-0.37605304),
            k = 16.03141911 + Tin / 100 * 1.74054886
        )
        return Fs * self.Fs_rated

def smooth_exp_transition(x, s, a, b, d,FUNC_where):
    """
    光滑连接的线性-指数分段函数
    
    参数：
    - x: 输入变量
    - s: 分段点
    - a: 线性部分斜率
    - b: 线性部分截距
    - d: 指数增长率
    """
    # 计算常数项
    c = a / (d * np.exp(d * s))
    e = a * s + b - a / d
    y = FUNC_where(
        x <= s,
        a * x + b,
        c * np.exp(d * x) + e
    )
    return y