"""
The module applies extinction and computes the apparent magntiude,  
when applied to data created by the code Galaxia.
It can also be used to compute proper motion and radial velocity, in 
equatorial coordinates.
Set GlobalGalaxiaPath  to correct value
It can also be used to run Galaxia from python by setting the 
Set GalaxiaExec to point to the executable file.

Copyright (c) 2019 Sanjib Sharma

NEW:
(Teff,logg) depdendent extinction for Gaia bands.
R_G= A_G/A_V

SETUP
::
In file gxutil.py provide the path to aebv_factor.ebf file.

USAGE
::
>>>import gxutil
>>>data=ebf.read(myfile,'/')
>>>gxutil.abs2app(data,corr=True)
>>>gxutil.append_pm(data)

"""

from __future__ import print_function
from __future__ import division
import ebf
import numpy as np
import scipy.interpolate
import os 
import os.path
import tempfile
import subprocess

GlobalGalaxiaPath='/work1/sharma/GalaxiaData/'
GalaxiaExec='GalaxiaN'        

def write_keyval(filename,data,keylist=None):
    """
    Write key value pairs from a file with first column being key and second value 
    Args:
        filename (str): name of file  
        data (dict):  
        keylist (sequence of str): the order of the keys in the file
    """
    with open(filename,'w') as fp:
        if keylist is None:
            keylist=list(data.keys())
        for key in keylist:
            s='%-32s %s \n'%(key, str(data[key]))
            fp.write(s)

def run_galaxia(params):
    galaxia=GalaxiaExec
    keylist=['outputFile','outputDir','photoSys','magcolorNames','appMagLimits[0]','appMagLimits[1]','absMagLimits[0]','absMagLimits[1]','colorLimits[0]','colorLimits[1]','geometryOption','surveyArea','fSample','popID','warpFlareOn','longitude','latitude','seed','r_max','starType','photoError','modelFile']
    fd,paramfile=tempfile.mkstemp(suffix='_galaxiaparam')
    os.close(fd)
    print(paramfile)
    write_keyval(paramfile,params,keylist)
    p1=subprocess.Popen([galaxia,'-r',paramfile])
    p1.wait()
    outfile1=params['outputDir']+params['outputFile']+'.ebf'
    os.remove(paramfile)



def galaxia_params(outfile,photosys,magcolornames,magmax,fsample=1.0,model='mrtd5'):
#    magcolornames='tmass_j,tmass_j-tmass_ks'
    param={}
    param['outputFile']=os.path.basename(outfile)
    param['outputDir']=os.path.dirname(outfile)
    if param['outputDir'] =='':
        param['outputDir']='./'
    param['photoSys']=photosys
    param['magcolorNames']=magcolornames
    param['appMagLimits[0]']=-1000.0
    param['appMagLimits[1]']=magmax
    param['absMagLimits[0]']=-1000.0
    param['absMagLimits[1]']=1000.0
    param['colorLimits[0]']=-1000.0
    param['colorLimits[1]']=1000
    param['geometryOption']=0
    param['surveyArea']=17044.7
    param['fSample']=fsample
    param['popID']=-1
    param['warpFlareOn']=1
    param['longitude']=10.0
    param['latitude']=90.0
    param['seed']=17
    param['r_max']=1000
    param['starType']=0
    param['photoError']=0
    param['modelFile']='Model/population_parameters_'+model+'.ebf'
    return param


def extcorr(ebv):
    return (1-np.tanh((ebv-0.15)/0.075))*0.5*(1-0.6)+0.6


def rename(data):
    data['log_teff']=data['teff'].copy()
    data['log_age']=data['age'].copy()
    data['log_lum']=data['lum'].copy()
    data['log_grav']=data['grav'].copy()
    del data['teff']
    del data['age']
    del data['lum']
    del data['grav']

    
def abs2app(data,noext=False,dered=False,corr=False,suffix='',rlambda=False,abmag=False):
    s=set(data.keys())
    aliaslist=['log_grav','log_teff','log_age']
    keylist=['grav','teff','age']
    for i,alias in enumerate(aliaslist):
        if alias not in s:
            data[alias]=data[keylist[i]]
            del data[keylist[i]]

    dmod=5.0*np.log10(100.0*data['rad'])
    if dered:
        ebv=data['exbv_schlegel']-data['exbv_schlegel_inf'] 
    else:
        ebv=data['exbv_schlegel']
    if corr:
        ebv=ebv*extcorr(data['exbv_schlegel_inf'])

#    aebv=aebv_factor()
    exto=Extinction(map3d=False)

    if abmag == True:
        for temp in list(data.keys()):
            if (temp.lower() in exto.aebv)or(('r_'+temp) in exto.r_lambda_data.keys()):
                data[temp+'_ab']=data[temp].copy()

    j=0L
    for temp in list(data.keys()):
        if (temp.lower() in exto.aebv)or(('r_'+temp) in exto.r_lambda_data.keys()):
            print(temp)
            if noext==True:
                data[temp+suffix]=data[temp]+dmod
            else: 
                if rlambda:
                    alambda=exto.r_lambda(data['log_teff'],data['log_grav'],temp.lower())*ebv
                    data[temp+suffix]=data[temp]+dmod+alambda
                else:
                    if temp.lower() in exto.aebv:
                        data[temp+suffix]=data[temp]+dmod+exto.aebv[temp.lower()]*ebv
                    else:
                        print('Deleting-',temp+suffix)
                        del data[temp+suffix]
            j+=1

    if  j == 0:
        raise RuntimeError('No quantity to add extinction')


def abs2app_file(infile,outfile,key=None,uplimit=14.0):
    data=ebf.read(infile,'/')
    abs2app(data,corr=True)
    print('Written:',outfile) 
    if key != None:
        ind=np.where(data[key]<uplimit)[0]
        for key1 in list(data.keys()):
            if data[key1].size == data[key].size:
                data[key1]=data[key1][ind]
    ebf.write(outfile,'/',data,'w')

def add_tmass_ph_error(data,fac=0.0):
    if 'tmass_ks' in data.keys():
        sig=0.01*np.exp(data['tmass_ks']-12.6)+0.0125+fac
    elif 'kepler_ks' in data.keys():
        sig=0.01*np.exp(data['kepler_ks']-12.6)+0.0125+fac
    elif 'dcmc_ks' in data.keys():
        sig=0.01*np.exp(data['dcmc_ks']-12.6)+0.0125+fac
    else:
        raise RuntimeError('dcmc_ks or kepler_ks not found')
    for key in ['dcmc_j','dcmc_h','dcmc_ks','dcmc_i','kepler_j','kepler_h','kepler_ks','tmass_j','tmass_h','tmass_ks']:
        if key in data.keys():
            data[key]=data[key]+np.random.normal(size=data[key].size)*sig



def lcolor(u,g,r,i):
    return -0.436*u+1.129*g-0.119*r-0.574*i+0.1984

def red_giant_logg(teff):
    status=0
    if hasattr(teff,'__len__') == False: 
        status=1
        temp=np.array([teff])
    temp=5.2-2.8e-4*teff
    ind=np.where(teff<4250.0)[0]
    if ind.size > 0:
        temp[ind]=4.0
    ind=np.where(teff>6000.0)[0]
    if ind.size > 0:
        temp[ind]=3.5
    if status:
        temp=temp[0] 
    return temp



def vxyz2lbr(px,py,pz,vx,vy,vz):
    r=np.sqrt(px*px+py*py+pz*pz)
    ind=np.where(r == 0.0)[0]
    r[ind]=1.0
    px=px/r
    py=py/r
    pz=pz/r
    rc=np.sqrt(px*px+py*py)
    ind=np.where(rc == 0.0)[0]
    rc[ind]=1.0
    tm_00=-py/rc; tm_01=-pz*px/rc; tm_02= rc*px/rc
    tm_10= px/rc; tm_11=-pz*py/rc; tm_12= rc*py/rc
    tm_20= 0.0  ; tm_21= rc  ; tm_22= pz
    vl=(vx*tm_00+vy*tm_10+vz*tm_20)
    vb=(vx*tm_01+vy*tm_11+vz*tm_21)
    vr=(vx*tm_02+vy*tm_12+vz*tm_22)
    return vl,vb,vr

def vxyz2lzr(px,py,pz,vx,vy,vz):
    r=np.sqrt(px*px+py*py+pz*pz)
    ind=np.where(r == 0.0)[0]
    r[ind]=1.0
    px=px/r
    py=py/r
    pz=pz/r
    rc=np.sqrt(px*px+py*py)
    ind=np.where(rc == 0.0)[0]
    rc[ind]=1.0
    tm_00=-py/rc; tm_01=0.0; tm_02= px/rc
    tm_10= px/rc; tm_11=0.0; tm_12= py/rc
    tm_20= 0.0  ; tm_21=1.0; tm_22= 0.0
    vl=(vx*tm_00+vy*tm_10+vz*tm_20)
    vz=(vx*tm_01+vy*tm_11+vz*tm_21)
    vr=(vx*tm_02+vy*tm_12+vz*tm_22)
    return vl,vz,vr


def xyz2lbr(x,y,z):
    rc2=x*x+y*y
    return [np.degrees(np.arctan2(y,x)),np.degrees(np.arctan(z/np.sqrt(rc2))),np.sqrt(rc2+z*z)]

def gal2equ(l, b,epoch='2000'):   
   g=ephem.Galactic(0.0,np.pi/2.0,epoch=str(epoch))
   [al_gp,de_gp]=g.to_radec()
   g.from_radec(0.0,np.pi/2.0)
   l_cp=g.lon.znorm   
   l = np.radians(l) 
   b = np.radians(b)

   al = al_gp + np.arctan2(np.cos(b) * np.sin(l_cp - l), np.cos(de_gp) * np.sin(b) - np.sin(de_gp) * np.cos(b) * np.cos(l_cp - l))   
   de=np.arcsin(np.sin(de_gp)*np.sin(b)+np.cos(de_gp)*np.cos(b)*np.cos(l_cp - l))

   de = np.degrees(de)      
   al = np.degrees(al)%360.0    
   return [al,de]


def gal2equ_pm(l, b, mul,mub,epoch='2000',distance=None):
    """
    l,b: degrees
    mul,mub:arcsec/year
    Dstance: kpc
    """
    g=ephem.Galactic(0.0,np.pi/2.0,epoch=str(epoch))
    [al_gp,de_gp]=g.to_radec()
    g.from_radec(0.0,np.pi/2.0)
    l_cp=g.lon.znorm   
    l = np.radians(l) 
    b = np.radians(b)

    al = al_gp + np.arctan2(np.cos(b) * np.sin(l_cp - l), np.cos(de_gp) * np.sin(b) - np.sin(de_gp) * np.cos(b) * np.cos(l_cp - l))   
    de=np.arcsin(np.sin(de_gp)*np.sin(b)+np.cos(de_gp)*np.cos(b)*np.cos(l_cp - l))
   
    c1=np.sin(de_gp)*np.cos(de)-np.cos(de_gp)*np.sin(de)*np.cos(al-al_gp)
    c2=np.cos(de_gp)*np.sin(al-al_gp)
    cosb=np.sqrt(c1*c1+c2*c2)
    mua=(c1*mul-c2*mub)/cosb
    mud=(c2*mul+c1*mub)/cosb   

    de = np.degrees(de)      
    al = np.degrees(al)%360.0    
    if distance is None:
        return [al,de,mua,mud]
    else:
        return [al,de,mua/(4.74e3*distance),mud/(4.74e3*distance)]



def append_muvr(data):
    """
    input in kpc and km/s
    d: dict with keys ['px','py','pz','log_age','pop_id']
    output
    pm in arcsec/year
    Adds: ['ra','dec',glon','glat','rad','vl','vb','vr','pmra','pmdec']
    """
    data['glon'],data['glat'],data['rad']=xyz2lbr(data['px'],data['py'],data['pz'])

    [vl,vb,vr]=vxyz2lbr(data['px'],data['py'],data['pz'],data['vx'],data['vy'],data['vz'])
    r=np.sqrt(data['px']*data['px']+data['py']*data['py']+data['pz']*data['pz'])
    data['pm']=np.sqrt(vl*vl+vb*vb)/(4.74e3*r)
    data['mu']=np.sqrt(vl*vl+vb*vb)/(4.74e3*r)
    data['ra'],data['dec'],data['pmra'],data['pmdec']=gal2equ_pm(data['glon'],data['glat'],vl/(4.74e3*r),vb/(4.74e3*r))
    data['vl']=vl
    data['vb']=vb
    data['vr']=vr


def append_radec(data):
    keys=list(data.keys())
    if ('glon' in keys) == False:
        data['glon'],data['glat'],r=autil.xyz2lbr(data['px'],data['py'],data['pz'])
    if ('ra' in keys) == False:
        data['ra'],data['dec']=autil.gal2equ(data['glon'],data['glat'])


def gal2ascii(infile,compress=True):
    outfile=infile.rsplit('.ebf',1)[0]+'.txt'
    data=ebf.read(infile,'/')
    nsize=data['px'].size
    keylist=['popid','satid','partid','fieldid','smass','mact','mtip','age','lum','teff','grav','feh','alpha','px','py','pz','vx','vy','vz','rad','ra','dec','glon','glat','exbv_schlegel_inf','exbv_schlegel','exbv_solar','mag0','mag1','mag2']

    dkeys=[key.lower() for key in list(data.keys())]
    skeys=[]
    for keyt in keylist:
        if data[keyt].size == nsize:
            if keyt in dkeys:
                skeys.append(keyt.lower())
    keys1=[]
    for keyt in list(data.keys()):
        if data[keyt].size == nsize:
            if keyt.lower() not in keylist:
                keys1.append(keyt.lower())
    for keyt in sorted(keys1):
        skeys.append(keyt)


    formatstring="{0:>20}"
    delimiter=','

    f=open(outfile,'w')
    f.write(delimiter.join(formatstring.format(key) for key in skeys)+'\n')
    for i in range(nsize):
        f.write(delimiter.join(formatstring.format(data[key][i]) for key in skeys)+'\n')
    f.close()

    if compress:
        ret=os.system("gzip "+outfile)




def mk_class(teff):
    edges=[0.0,3700.0,5200.0,6000.0 ,7500.0,1e4,3e4,1e10]
    #    mktype=['M','K','G','F','A','B','O']
    #    mktype=['0','1','2','3','4','5','6']
    if (hasattr(teff,"__len__") == False):
        return np.digitize([teff],bins=edges)[0]-1
    else:
        return np.digitize(teff,bins=edges)-1

def mk_class_hist(teff):
    edges=[0.0,3700.0,5200.0,6000.0,7500.0,1e4,3e4,1e10]
    #    mktype=['M','K','G','F','A','B','O']
    #    mktype=['0','1','2','3','4','5','6']
    return np.histogram(teff,bins=edges)[0][::-1]

class Extinction(object):
    def __init__(self,map3d=True):
        # Extinction 
#        self.GalaxiaPath='/work1/sharma/GsynthData/'        
        self.GalaxiaPath=GlobalGalaxiaPath  
        GalaxiaPath1=self.GalaxiaPath+'Extinction/'
        if map3d:
            data3d=ebf.read(GalaxiaPath1+'ExMap3d_1024.ebf','/ExMap3d.data')
            xmms3d=ebf.read(GalaxiaPath1+'ExMap3d_1024.ebf','/ExMap3d.xmms')
            points3d=[xmms3d[i,0]+np.arange(data3d.shape[i])*xmms3d[i,2] for i in range(data3d.ndim)]
            data2d=ebf.read(GalaxiaPath1+'Schlegel_4096.ebf','/ExMap2d.data')
            xmms2d=ebf.read(GalaxiaPath1+'Schlegel_4096.ebf','/ExMap2d.xmms')
            points2d=[xmms2d[i,0]+np.arange(data2d.shape[i])*xmms2d[i,2] for i in range(data2d.ndim)]
            self.ebv3d = scipy.interpolate.RegularGridInterpolator(points3d, data3d,bounds_error=False,fill_value=None,method='linear')
            self.ebv2d = scipy.interpolate.RegularGridInterpolator(points2d, data2d,bounds_error=False,fill_value=None,method='linear')

        self.aebv=self.aebv_factor()
        self.r_lambda_data=ebf.read(self.GalaxiaPath+'Isochrones/r_lambda_parsec.ebf','/')

    def aebv_factor(self):
        filename=self.GalaxiaPath+'Isochrones/aebv_factor.ebf'
        aebvdata=ebf.read(filename,'/')
        aebv={}
        for myfilter,factor in zip(aebvdata['filter'],aebvdata['aebv_factor']):
            aebv[str(myfilter).lower()]=factor
        return aebv

    def ebv(self,glon,glat,dist,corr=False):
        glon=np.array(glon)%360.0
        x=np.zeros((glon.size,3),dtype='float64')
        x[:,0]=glon
        x[:,1]=glat
        x[:,2]=np.log10(dist)
        temp3d=self.ebv3d(x)
        temp2d=self.ebv2d(x[:,0:2])
        if corr:
            temp2d=temp2d*extcorr(temp2d)
        if glon.ndim == 0:
            return temp2d[0]*temp3d[0]
        else:
            return temp2d*temp3d

    def r_lambda(self,log_teff,log_grav,bandname,simple=False):
#        s=set(['gaia_g','gaia_gbp','gaia_grp','tmass_j','tmass_h','tmass_ks'])
        if simple:
            return self.aebv[bandname]+log_teff-log_teff

        rg=self.r_lambda_data
        if 'r_'+bandname in list(rg.keys()):
            indd=np.where((rg['log_grav']>4.0)|(rg['log_teff']>np.log10(15000)))[0]
            indg=np.where((rg['log_grav']<3.0)|(rg['log_teff']>np.log10(15000)))[0]
            funcd=scipy.interpolate.interp1d(rg['log_teff'][indd],rg['r_'+bandname][indd],fill_value='extrapolate',assume_sorted=True)
            funcg=scipy.interpolate.interp1d(rg['log_teff'][indg],rg['r_'+bandname][indg],fill_value='extrapolate',assume_sorted=True)
            loc=3.0
            scale=0.1
            y1=(1+np.tanh((log_grav-loc)/scale))*0.5*funcd(log_teff)
            y2=(1-np.tanh((log_grav-loc)/scale))*0.5*funcg(log_teff)
            return y1+y2
        else:
            return self.aebv[bandname]+log_teff-log_teff


    def redden(self,ebv,mag,bandname):
        return mag+ebv*self.aebv[bandname]

    def deredden(self,ebv,mag,bandname):
        return mag-ebv*self.aebv[bandname]

    def ebv_inf(self,glon,glat):
        glon=np.array(glon)
        glat=np.array(glat)
        x=np.zeros((glon.size,2),dtype='float64')
        x[:,0]=glon
        x[:,1]=glat
        if glon.ndim == 0:
            return self.ebv2d(x)[0]
        else:
            return self.ebv2d(x)