###############################################################################
# gaia_rc.py: functions for the RC in Gaia
###############################################################################
import numpy
from scipy import optimize
import isodist
from apogee.samples import rc
from apogee.util import localfehdist
def G_jordi(g,gz):
    """G in terms of g and g-z From Jordi et al. (2010)"""
    return g-0.1154-0.4175*gz-0.0497*gz**2.+0.0016*gz**3. 
def G(age,Z,iso):
    """G for a given age and metallicity"""
    Zs= iso.Zs()
    tZ= Zs[numpy.argmin(numpy.fabs(Z-Zs))]
    p= iso(age,tZ)
    jk= p['J']-p['Ks']
    indx= (jk < 0.8)*(jk > 0.5)\
        *(tZ <= 0.06)\
        *(tZ <= rc.jkzcut(jk,upper=True))\
        *(tZ >= rc.jkzcut(jk))\
        *(p['logg'] >= rc.loggteffcut(10.**p['logTe'],tZ,upper=False))\
        *(p['logg'] <= rc.loggteffcut(10.**p['logTe'],tZ,upper=True))
    outG= G_jordi(p['g'],p['g']-p['z'])
    # Average over the IMF
    sindx= numpy.argsort(p['M_ini'])
    outG= outG[sindx]
    int_IMF= p['int_IMF'][sindx]
    w= (numpy.roll(int_IMF,-1)-int_IMF)/(int_IMF[-1]-int_IMF[0])
    w= w[indx]
    outG= outG[indx]
    return (numpy.nansum(w*outG)/numpy.nansum(w),
            numpy.sqrt(numpy.nansum(w*outG**2.)/numpy.nansum(w)
                       -(numpy.nansum(w*outG)/numpy.nansum(w))**2.))
def Gdist(tG,iso):
    """The distribution of G for the local metallicity distribution
    REWRITE TO TAKE A Z(G) FUNCTION OBTAINED THROUGH INTERPOLATION"""
    try:
        tZ= optimize.brenth(lambda x: G(9.7,x,iso)[0]-tG,0.0001,0.06)
    except ValueError:
        return 0.
    # Add Jacobian
    return localfehdist(isodist.Z2FEH(tZ,zsolar=0.017))*1.
def load_iso():
    Zs= [0.0035,0.017,0.035]
    return isodist.PadovaIsochrone(type='sdss-2mass',parsec=True,
                                   Z=Zs)

@@ -227,23 +227,26 @@
  import 
-optimiz
+interpolat
 e%0Aimport
@@ -1351,19 +1351,18 @@
 dist(tG,
-iso
+ZG
 ):%0A    %22
@@ -1519,64 +1519,85 @@
 tZ= 
-optimize.brenth(lambda x: G(9.7,x,iso)%5B0%5D-tG,0.0001,0.06
+ZG(tG)%0A        tjac= numpy.fabs(ZG.derivatives(tG)%5B0%5D*0.017/tZ/numpy.log(10.)
 )%0A  
@@ -1714,10 +1714,290 @@
 7))*
-1.
+tjac%0Adef sample_Gdist(iso,n=1000):%0A    %22%22%22Sample from the distribution of MG%22%22%22%0A    # First calculate the ditribution%0A    Gs= numpy.linspace(0.1,1.9,201)%0A    ZG= load_ZG(iso)%0A    pG= numpy.array(%5BGdist(g,ZG) for g in Gs%5D)%0A    pGmax= numpy.nanmax(pG)%0A    # Now rejection sample%0A    %0A
 %0Adef
@@ -2021,28 +2021,42 @@
 Zs= 
-%5B
+numpy.arange(
 0.00
-3
+0
 5,0.0
-17
+605
 ,0.0
-35%5D
+005)
 %0A   
@@ -2159,8 +2159,431 @@
   Z=Zs)%0A
+def load_ZG(iso,s=0.00005):%0A    %22%22%22Return a function that gives Z as a function of G%22%22%22%0A    Zs= iso.Zs()%0A    Zs= Zs%5BZs %3C 0.05%5D%0A    tage= iso.logages()%5B62%5D%0A    Gs= numpy.array(%5BG(tage,z,iso)%5B0%5D for z in Zs%5D)%0A    sindx= numpy.argsort(Gs)%0A    Zs= Zs%5Bsindx%5D%0A    Gs= Gs%5Bsindx%5D%0A    goodIndx= True-numpy.isnan(Gs)%0A    return interpolate.UnivariateSpline(Gs%5BgoodIndx%5D,Zs%5BgoodIndx%5D,k=3,%0A                                        s=s)%0A

determine p(G) from local feh