from __future__ import print_function
import numpy as np


def sigma_lz(x,lz):
    return (x[3]*np.power(lz/(8*241.92),2)+np.exp(-(lz-8*241.92)/x[2]))/(1+x[3])

def sigma_age(x,age):
    return np.power((age+0.1)/(10.1),x[1])

def sigma_feh(x,fe_h):
    return (1+x[4]*fe_h)

def sigma_z(x,pz):
    return 1+np.abs(pz)*x[5]

def velocity_dispersion(lz,age,fe_h,pz,option='vz'):
    """
    Arguments
    ---------
        lz: (float) or None,   R*v_phi [kpc km/s]
        age: (float) or None[Gyr]
        fe_h: (float) or None  
        pz: (float) or None [kpc], height from midplane
 
    Keywords
    --------
        option: (str) 'vz' or 'vr'
    
     
    Examples:
    ------------------
    velocity_dispersion(None,np.array([1,2,3]),None,None)
    """
    if option == 'vz':
        x=np.array([21.1, 0.441, 1130, 0.58, -0.52, 0.20, 1/(20*2.25)])
    elif option == 'vr':
        x=np.array([39.4, 0.251, 2300, 0.09, -0.19, 0.12, 1/(40*2.25)])        
    else:
        raise RuntimeError('option should be vz or vr')
    d={'age':age,'lz':lz,'fe_h':fe_h,'pz':pz}
    sigma=x[0]
    for key in ['lz','age','fe_h','pz']:        
        if d[key] is not None:
                if key == 'lz':
                sigma=sigma*sigma_lz(x,d['lz'])
            elif key == 'age':
                sigma=sigma*sigma_age(x,d['age'])
            elif key == 'fe_h':
                sigma=sigma*sigma_feh(x,d['fe_h'])
            elif key == 'pz':
                sigma=sigma*sigma_z(x,d['pz'])
return sigma