import ebf3 as ebf
import numpy as np
import scipy.stats

"""
import abundmodel_pub
am_msto=abundmodel_pub(dlabel='msto')
am_giant=abundmodel_pub(dlabel='giant')
age=np.linspace(0,13,100)
plt.plot(age,am.abund('ca_fe',age,0.0))
feh=np.linspace(-2,0.5,100)
plt.plot(feh,am.abund('ca_fe',4.6,fe_h))

"""

class AbundModel:
    def __init__(self,dlabel='msto'):
        self.d=ebf.read('abundmodel_compact.ebf','/'+dlabel+'/',recon=1)
        self.elements=['fe_h','alpha_fe','mg_fe','si_fe', 'ca_fe','ti2_fe', 'ti_fe', 'o_fe', 'c_fe', 'sc_fe','k_fe','na_fe', 'cu_fe', 'mn_fe','cr_fe', 'ni_fe',  'ba_fe','y_fe','ce_fe', 'ru_fe',  'al_fe','zn_fe','nd_fe','eu_fe',  'co_fe', 'la_fe', 'li_fe', 'v_fe', 'rb_fe', 'sm_fe', 'mo_fe', 'zr_fe','sr_fe']
        self.labels=['Fe',r'$\alpha$','Mg','Si','Ca','Ti2', 'Ti', 'O','C','Sc','K','Na','Cu','Mn','Cr','Ni','Ba','Y', 'Ce','Ru', 'Al','Zn','Nd','Eu','Co','La','Li','V','Rb','Sm','Mo','Zr','Sr'] 


    def element_symbol(self,x):
        if x == 'alpha_fe':
            return r'$\alpha$'
        else:
            return x.split('_')[0].capitalize()
        
    def abund(self,key,age,feh):
        """
        Get abundace of an element as function of age andf [Fe/H]
        Arguments
        ----------
        key: name of element whose abundace is desired
        age: Age [Gyr]
        feh: [Fe/H] [dex]
        """
        f_age=self.abund_age(key,age)
        f_fe_h=self.abund_fe_h(key,feh)
        return (f_fe_h+f_age)+self.d[key]['solar_offset1']

    def abund_age(self,key,age):
        """
        Get abundace of an element as function of age for solar metallicity
        Arguments
        ----------
        key: name of element whose abundace is desired
        age: Age [Gyr]
        """
        return np.interp(age,self.d[key]['age'],self.d[key]['abund_age'])

    def abund_fe_h(self,key,feh):
        """
        Get abundace of an element as function of [Fe/H] for solar age 
        Arguments
        ----------
        key: name of element whose abundace is desired
        feh: [Fe/H] [dex]
        """
        return np.interp(feh,self.d[key]['fe_h'],self.d[key]['abund_fe_h'])


    def pdf_age(self,key,age,fe_h,e_fe_h,abund,e_abund):
        """
        Probability that for a given age and [Fe/H] the abundace of element 
        "key" is abund. 
        Arguments
        ----------
        key: name of element
        age: Age [Gyr]
        feh: [Fe/H] [dex]
        e_feh: error in [Fe/H]
        abund: abundace of "key"
        e_abund: error in abundance
        """
        return scipy.stats.norm.pdf(abund-self.abund(key,age,fe_h),scale=e_abund)