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
    ):
        Fs = (ToutA - TinA) * FA
        return {'Fs':Fs}
    
    def prior(self):
        param = {}
        return param


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