import numpy as np
import os
from datetime import datetime
from scipy.interpolate import InterpolatedUnivariateSpline, UnivariateSpline, interp1d
from scipy import interpolate
from astropy.table import Table
import galsim

vc_A = 2.99792458e+18  # speed of light: A/s
vc_m = 2.99792458e+8   # speed of light: m/s
h_Plank = 6.626196e-27 # Plank constant: erg s

def magToFlux(mag):
    """
    flux of a given AB magnitude

    Parameters:
    mag: magnitude in unit of AB

    Return:
    flux: flux in unit of erg/s/cm^2/Hz
    """
    flux = 10**(-0.4*(mag+48.6))
    return flux

def extAv(nav, seed=1212123):
    """
    Generate random intrinsic extinction Av
    following the distribution from Holwerda et al, 2015
    """
    np.random.seed(seed)
    tau = 0.4
    peak, a = 0.1, 0.5
    b = a*(tau-peak)
    pav = lambda av: (a*av+b)*np.exp(-av/tau)
    avmin, avmax = 0., 3.
    avs = np.linspace(avmin, avmax, int((avmax-avmin)/0.001)+1)
    norm = np.trapz(pav(avs), avs)
    pav_base = pav(avs)/np.sum(pav(avs))
    rav = np.random.choice(avs, nav, p=pav_base)
    return rav


###########################################
def seds(sedlistn, seddir="./", unit="A"):
    """
    read SEDs and save into Python dictionary

    Parameters:
    sedlistn: filename of the sed template list and corresponding intrinsic 
              extinction model, see tmag_dz.py for detailes
    listdir:  directory of the list
    unit: wavelength unit of the input templates

    Return:
    SED dictionary, the output wavelength unit is 'A'
    """
    seds = {}
    reds = {}
    sedlist = seddir + sedlistn
    sedn = open(sedlist).read().splitlines()
    sedtype = range(1,len(sedn)+1)
    for i in range(len(sedn)):
        xxx = sedn[i].split()
        isedn = seddir+xxx[0]
        itype = sedtype[i]
        ised = np.loadtxt(isedn)
        if unit=="nm": ised[:,0] *= 10.0
        seds[itype] = ised
        reds[itype] = int(xxx[1])

    return seds, reds

def sed_assign(phz, btt, rng):
    """
    assign SED template to a galaxy.
    """
    sedid = list(range(1, 34))
    lzid = sedid[0:13] + sedid[23:28]
    hzid = sedid[13:23]
    if btt == 1.0:
        sedtype = rng.sample(sedid[0:5] + sedid[28:33], 1)[0]
        if phz > 2.0 and sedtype in sedid[0:5]:
            sedtype = rng.sample(sedid[28:33], 1)[0]
    elif btt > 0.3 and btt < 1.0:
        sedtype = rng.sample(sedid, 1)[0]
        if phz > 2.0 and sedtype in lzid:
            sedtype = rng.sample(hzid, 1)[0]
    elif btt >= 0.1 and btt < 0.3:
        sedtype = rng.sample(sedid[5:28], 1)[0]
        if phz > 1.5 and sedtype in lzid:
            sedtype = rng.sample(hzid, 1)[0]
    elif btt >= 0.0 and btt < 0.1:
        sedtype = rng.sample(sedid[5:23], 1)[0]
        if phz > 1.5 and sedtype in lzid:
            sedtype = rng.sample(hzid, 1)[0]
    else:
        sedtype = 0
    return sedtype

###########################################
def tflux(filt, sed, redshift=0.0, av=0.0, redden=0):
    """
    calculate the theoretical SED for given filter set and template
    Only AB magnitude is support!!!

    Parameters:
    filt: 2d array
        fliter transmissions: lambda(A), T
    sed: 2d array
        sed templateL lambda (A), flux (erg s-1 cm-2 A-1)
    redshift: float
        redshift of the corresponding source, default is zero
    av: float
        extinction value for intrincal extinction
    redden: int
        reddening model, see Function 'reddening' for details

    return:
    tflux: float
        theoretical flux
    sedObs: array
        SED in observed frame
    """
    z = redshift + 1.0
    sw, sf = sed[:,0], sed[:,1]
    # reddening
    sf = reddening(sw, sf, av=av, model=redden)
    sw, sf = sw*z, sf*(z**3)
    # lyman forest correction
    sf = lyman_forest(sw, sf, redshift)
    
    sedxx = (sw.copy(), sf.copy())
    
    sw = vc_A/sw
    sf = sf*(vc_A/sw**2) # convert flux unit to erg/s/cm^s/Hz
    sw, sf = sw[::-1], sf[::-1]
    sfun = interp1d(sw, sf, kind='linear')

    fwave, fresp = filt[:,0], filt[:,1]
    fwave = vc_A/fwave
    fwave, fresp = fwave[::-1], fresp[::-1]
    tflux = sfun(fwave)

    zpflux = 3.631*1.0e-20

    tflux = np.trapz(tflux*fresp/fwave,fwave)/np.trapz(zpflux*fresp/fwave,fwave)
    #tflux = np.trapz(tflux*fresp,fwave)/np.trapz(zpflux*fresp,fwave)

    return tflux, sedxx

###########################################
def lyman_forest(wavelen, flux, z):
    """
    Compute the Lyman forest mean absorption of an input spectrum,
    according to D_A and D_B evolution from Madau (1995).
    The waveeln and flux are in observed frame
    """
    if z<=0:
        flux0 = flux
    else:
        nw = 200
        istep = np.linspace(0,nw-1,nw)
        w1a, w2a = 1050.0*(1.0+z), 1170.0*(1.0+z)
        w1b, w2b = 920.0*(1.0+z), 1015.0*(1.0+z)
        wstepa = (w2a-w1a)/float(nw)
        wstepb = (w2b-w1b)/float(nw)

        wtempa = w1a + istep*wstepa
        ptaua = np.exp(-3.6e-03*(wtempa/1216.0)**3.46)

        wtempb = w1b + istep*wstepb
        ptaub = np.exp(-1.7e-3*(wtempb/1026.0)**3.46\
                       -1.2e-3*(wtempb/972.50)**3.46\
                       -9.3e-4*(wtempb/950.00)**3.46)

        da = (1.0/(120.0*(1.0+z)))*np.trapz(ptaua, wtempa)
        db = (1.0/(95.0*(1.0+z)))*np.trapz(ptaub, wtempb)

        if da>1.0: da=1.0
        if db>1.0: db=1.0
        if da<0.0: da=0.0
        if db<0.0: db=0.0
        flux0 = flux.copy()
        id0 = wavelen<=1026.0*(1.0+z)
        id1 = np.logical_and(wavelen<1216.0*(1.0+z),wavelen>=1026.0*(1.0+z))
        flux0[id0] = db*flux[id0]
        flux0[id1] = da*flux[id1]

    return flux0


###########################################
def reddening(sw, sf, av=0.0, model=0):
    """
    calculate the intrinsic extinction of a given template

    Parameters:
    sw: array
        wavelength
    sf: array
        flux
    av: float or array
    model: int
        Five models will be used:
        1: Allen (1976) for the Milky Way
        2: Seaton (1979) fit by Fitzpatrick (1986) for the Milky Way
        3: Fitzpatrick (1986) for the Large Magellanic Cloud (LMC)
        4: Prevot et al (1984) and Bouchet (1985) for the Small Magellanic Cloud (SMC)
        5: Calzetti et al (2000) for starburst galaxies
        6: Reddy et al (2015) for star forming galaxies

    Return:
    reddening-corrected flux or observed flux
    """
    if model==0 or av==0.0:
        flux=sf
    elif model==1: # Allen (1976) for the Milky Way
        lambda0 = np.array([1000, 1110, 1250, 1430, 1670, \
                            2000, 2220, 2500, 2850, 3330, \
                            3650, 4000, 4400, 5000, 5530, \
                            6700, 9000, 10000, 20000, 100000], dtype=float)
        kR = np.array([4.20, 3.70, 3.30, 3.00, 2.70, \
                       2.80, 2.90, 2.30, 1.97, 1.69, \
                       1.58, 1.45, 1.32, 1.13, 1.00, \
                       0.74, 0.46, 0.38, 0.11, 0.00],dtype=float)
        ext0 = InterpolatedUnivariateSpline(lambda0, kR, k=1)
        A_lambda = av*ext0(sw)
        A_lambda[A_lambda<0.0] = 0.0
        flux = sf*10**(-0.4*A_lambda)
    elif model==2: # Seaton (1979) fit by Fitzpatrick (1986) for the Milky Way
        Rv=3.1
        al0, ga, c1, c2, c3, c4 = 4.595, 1.051, -0.38, 0.74, 3.96, 0.26
        ff11 = __red(1100.0,al0,ga,c1,c2,c3,c4)
        ff12 = __red(1200.0,al0,ga,c1,c2,c3,c4)
        slope=(ff12-ff11)/100.0
        lambda0 = np.array([3650, 4000, 4400, 5000, 5530, \
                            6700, 9000, 10000, 20000, 100000], dtype=float)
        kR = np.array([1.58, 1.45, 1.32, 1.13, 1.00, \
                       0.74, 0.46, 0.38, 0.11, 0.00],dtype=float)
        fun = interp1d(lambda0, kR, kind='linear')

        sw0 = sw[sw<1200.0]
        A_lambda0 = (ff11+(sw0-1100.0)*slope)/Rv+1.0
        sw1 = sw[np.logical_and(sw>=1200.0, sw<=3650.0)]
        ff = __red(sw1,al0,ga,c1,c2,c3,c4)
        A_lambda1 = ff/Rv+1.0
        sw2 = sw[np.logical_and(sw>3650.0, sw<=100000.0)]
        A_lambda2 = fun(sw2)
        A_lambda3 = sw[sw>100000.0]*0.0
        A_lambda = av*np.hstack([A_lambda0,A_lambda1,A_lambda2,A_lambda3])
        A_lambda[A_lambda<0.0] = 0.0
        flux = sf*10**(-0.4*A_lambda)
    elif model==3: # Fitzpatrick (1986) for the Large Magellanic Cloud (LMC)
        Rv=3.1
        al0, ga, c1, c2, c3, c4 = 4.608, 0.994, -0.69, 0.89, 2.55, 0.50
        ff11 = __red(1100.0,al0,ga,c1,c2,c3,c4)
        ff12 = __red(1200.0,al0,ga,c1,c2,c3,c4)
        slope=(ff12-ff11)/100.0
        lambda0 = np.array([3330, 3650, 4000, 4400, 5000, 5530, \
                            6700, 9000, 10000, 20000, 100000], dtype=float)
        kR = np.array([1.682, 1.58, 1.45, 1.32, 1.13, 1.00, \
                       0.74, 0.46, 0.38, 0.11, 0.00],dtype=float)
        fun = interp1d(lambda0, kR, kind='linear')

        sw0 = sw[sw<1200.0]
        A_lambda0 = (ff11+(sw0-1100.0)*slope)/Rv+1.0
        sw1 = sw[np.logical_and(sw>=1200.0, sw<=3330.0)]
        ff = __red(sw1,al0,ga,c1,c2,c3,c4)
        A_lambda1 = ff/Rv+1.0
        sw2 = sw[np.logical_and(sw>3330.0, sw<=100000.0)]
        A_lambda2 = fun(sw2)
        A_lambda3 = sw[sw>100000.0]*0.0
        A_lambda = av*np.hstack([A_lambda0,A_lambda1,A_lambda2,A_lambda3])
        A_lambda[A_lambda<0.0] = 0.0
        flux = sf*10**(-0.4*A_lambda)
    elif model==4: # Prevot et al (1984) and Bouchet (1985) for the Small Magellanic Cloud (SMC)
        Rv = 2.72
        lambda0 = np.array([1275, 1330, 1385, 1435, 1490, 1545, \
                            1595, 1647, 1700, 1755, 1810, 1860, \
                            1910, 2000, 2115, 2220, 2335, 2445, \
                            2550, 2665, 2778, 2890, 2995, 3105, \
                            3704, 4255, 5291, 12500, 16500, 22000], dtype=float)
        kR = np.array([13.54, 12.52, 11.51, 10.80, 9.84, 9.28, \
                        9.06, 8.49, 8.01, 7.71, 7.17, 6.90, 6.76, \
                        6.38, 5.85, 5.30, 4.53, 4.24, 3.91, 3.49, \
                        3.15, 3.00, 2.65, 2.29, 1.81, 1.00, 0.00, \
                        -2.02, -2.36, -2.47],dtype=float)
        kR = kR/Rv+1.0
        ext0 = InterpolatedUnivariateSpline(lambda0, kR, k=1)
        A_lambda = av*ext0(sw)
        A_lambda[A_lambda<0.0] = 0.0
        flux = sf*10**(-0.4*A_lambda)
    elif model==5: # Calzetti et al (2000) for starburst galaxies
        Rv = 4.05
        sw = sw*1.0e-04           #wavelength in microns

        fun1 = lambda x: 2.659*(-2.156+1.509/x-0.198/x**2+0.011/x**3)+Rv
        fun2 = lambda x: 2.659*(-1.857+1.040/x)+Rv

        ff11, ff12 = fun1(0.11), fun1(0.12)
        slope1=(ff12-ff11)/0.01
        ff99, ff100 = fun2(2.19), fun2(2.2)
        slope2=(ff100-ff99)/0.01

        sw0 = sw[sw<0.12]
        sw1 = sw[np.logical_and(sw>=0.12, sw<=0.63)]
        sw2 = sw[np.logical_and(sw>0.63, sw<=2.2)]
        sw3 = sw[sw>2.2]
        k_lambda0 = ff11+(sw0-0.11)*slope1
        k_lambda1, k_lambda2 = fun1(sw1), fun2(sw2)
        k_lambda3 = ff99+(sw3-2.19)*slope2
        A_lambda = av*np.hstack([k_lambda0,k_lambda1,k_lambda2,k_lambda3])/Rv
        A_lambda[A_lambda<0.0] = 0.0
        flux = sf*10**(-0.4*A_lambda)
    elif model==6: # Reddy et al (2015) for satr forming galaxies
        Rv = 2.505
        sw = sw*1.0e-04

        fun1 = lambda x: -5.726+4.004/x-0.525/x**2+0.029/x**3+Rv
        fun2 = lambda x: -2.672-0.010/x+1.532/x**2-0.412/x**3+Rv

        ff11, ff12 = fun1(0.14), fun1(0.15)
        slope1=(ff12-ff11)/0.01
        ff99, ff100 = fun2(2.84), fun2(2.85)
        slope2=(ff100-ff99)/0.01

        sw0 = sw[sw<0.15]
        sw1 = sw[np.logical_and(sw>=0.15, sw<0.60)]
        sw2 = sw[np.logical_and(sw>=0.60, sw<2.85)]
        sw3 = sw[sw>=2.85]
        k_lambda0 = ff11+(sw0-0.14)*slope1
        k_lambda1, k_lambda2 = fun1(sw1), fun2(sw2)
        k_lambda3 = ff99+(sw3-2.84)*slope2
        A_lambda = av*np.hstack([k_lambda0,k_lambda1,k_lambda2,k_lambda3])/Rv
        A_lambda[A_lambda<0.0] = 0.0
        flux  = sf*10**(-0.4*A_lambda)

    else:
        print("!!! Please select a proper reddening model")
        sys.exit(0)

    return flux

###########################################
def __red(alan,al0,ga,c1,c2,c3,c4):

    fun1 = lambda x: c3/(((x-(al0**2/x))**2)+ga*ga)
    fun2 = lambda x,cc: cc*(0.539*((x-5.9)**2)+0.0564*((x-5.9)**3))
    fun = lambda x,cc: c1+c2*x+fun1(x)+fun2(x,cc)

    ala = alan*1.0e-04 #wavelength in microns
    p = 1.0/ala
    if np.size(p)>1:
        p1, p2 = p[p>=5.9], p[p<5.9]
        ff = np.append(fun(p1,c4), fun(p2,0.0))
    elif np.size(p)==1:
        if p<5.9: c4 = 0.0
        ff = fun(p, c4)
    else:
        return
    return ff

###########################################
def sed2mag(mag_i, sedCat, filter_list, redshift=0.0, av=0.0, redden=0):

    # load the filters
    nfilt = len(filter_list)
    aflux = np.zeros(nfilt)
    nid = -1
    for k in range(nfilt):
        if filter_list[k].filter_type == 'i':
            nid = k
        bandpass = filter_list[k].bandpass_full
        ktrans = np.transpose(np.array([bandpass.wave_list*10.0, bandpass.func(bandpass.wave_list)]))
        aflux[k], isedObs = tflux(ktrans, sedCat, redshift=redshift, av=av, redden=redden)

    # normalize to i-band
    aflux = aflux / aflux[nid]

    # magnitudes in all filters
    amag = -2.5*np.log10(aflux) + mag_i
    spec = galsim.LookupTable(x=np.array(isedObs[0]), f=np.array(isedObs[1]), interpolant='nearest')
    isedObs = galsim.SED(spec, wave_type='A', flux_type='1', fast=False)
    return amag, isedObs

def eObs(e1,e2,g1,g2):
    """
    Calculate the sheared (observed) ellipticity using the 
    intrinsic ellipticity and cosmic shear components.

    Parameters:
    e1, e2: scalar or numpy array
    g1, g2: scalar or numpy array

    Return:
    Sheared (observed) ellipticity components, absolute value, and orientation
    in format of scalar or numpy array

    ** NOTE: e1, e2, g1, and g2 should have the same dimensions.
    """
    if np.isscalar(e1):
        e1 = np.array([e1])
        e2 = np.array([e2])
        g1 = np.array([g1])
        g2 = np.array([g2])
    else:
        e1 = e1.flatten()
        e2 = e2.flatten()
        g1 = g1.flatten()
        g2 = g2.flatten()

    # calculate the sheared (observed) ellipticity using complex rule
    nobj = len(e1)
    e1obs, e2obs = [], []
    eeobs, theta = [], []
    for i in range(nobj):
        e = complex(e1[i], e2[i])
        g = complex(g1[i], g2[i])
        e, gg = abs(e), abs(g)
        if gg<=1.0:
            tt = e + g
            bb = 1.0 + e*g.conjugate()
            eobs = tt/bb
        else:
            tt = 1.0 + g*e.conjugate()
            bb = e.conjugate() + g.conjugate()
            eobs = tt/bb
    
    # derive the orientation
        dd = 0.5*np.arctan(abs(eobs.imag/eobs.real))*180.0/np.pi
        if eobs.imag>0 and eobs.real>0: dd = dd
        if eobs.imag>0 and eobs.real<0: dd = 90.0 - dd
        if eobs.imag<0 and eobs.real>0: dd =  0.0 - dd
        if eobs.imag<0 and eobs.real<0: dd = dd - 90.0

        e1obs += [eobs.real]
        e2obs += [eobs.imag]
        eeobs += [abs(eobs)]
        theta += [dd]

    e1obs,e2obs,eeobs,theta = np.array(e1obs),np.array(e2obs),np.array(eeobs),np.array(theta)
    if nobj == 1: e1obs,e2obs,eeobs,theta = e1obs[0],e2obs[0],eeobs[0],theta[0]
    
    return e1obs, e2obs, eeobs, theta

def getObservedSED(sedCat, redshift=0.0, av=0.0, redden=0):
    z = redshift + 1.0
    sw, sf = sedCat[:,0], sedCat[:,1]
    # reddening
    sf = reddening(sw, sf, av=av, model=redden)
    sw, sf = sw*z, sf*(z**3)
    # lyman forest correction
    sf = lyman_forest(sw, sf, redshift)
    isedObs = (sw.copy(), sf.copy())
    return isedObs

def integrate_sed_bandpass(sed, bandpass):
    wave = np.linspace(bandpass.blue_limit, bandpass.red_limit, 1000) # in nm
    flux_normalized = sed(wave)*bandpass(wave)
    # print('in integrate_sed_bandpass', bandpass.blue_limit, bandpass.red_limit)
    int_flux = np.trapz(y=flux_normalized, x=wave) * 10. # convert to photons s-1 m-2 A-1
    return int_flux

def getABMAG(interFlux, bandpass):
    throughtput = Table(np.array(np.array([bandpass.wave_list*10.0, bandpass.func(bandpass.wave_list)])).T, names=(['WAVELENGTH', 'SENSITIVITY']))
    sWave = bandpass.blue_limit*10.0
    eWave = bandpass.red_limit*10.0
    # print('in getABMAG', sWave, eWave)
    ABMAG_zero = getABMagAverageVal(
        ABmag=0, 
        norm_thr=throughtput, 
        sWave=sWave, 
        eWave=eWave)
    flux_ave = interFlux / (eWave-sWave)
    ABMAG_spec = -2.5 * np.log10(flux_ave/ABMAG_zero)
    return ABMAG_spec

def getABMagAverageVal(ABmag=20.,norm_thr=None, sWave=6840, eWave=8250):
    """
    norm_thr: astropy.table, 2 colum, 'WAVELENGTH', 'SENSITIVITY'

    Return:
        the integerate flux of AB magnitude in the norm_thr(the throughtput of band), photos s-1 m-2 A-1
    """

    inverseLambda = norm_thr['SENSITIVITY']/norm_thr['WAVELENGTH']
    norm_thr_i = interpolate.interp1d(norm_thr['WAVELENGTH'], inverseLambda)

    x = np.linspace(sWave,eWave, int(eWave)-int(sWave)+1)
    y = norm_thr_i(x)
    AverageLamdaInverse = np.trapz(y,x)/(eWave-sWave)
    norm = 54798696332.52474 * pow(10.0, -0.4 * ABmag) * AverageLamdaInverse
    # print('AverageLamdaInverse = ', AverageLamdaInverse)
    # print('norm = ', norm)

    return norm

def getNormFactorForSpecWithABMAG(ABMag=20., spectrum=None, norm_thr=None, sWave=6840, eWave=8250):
    """
    Use AB magnitude system (zero point, fv = 3631 janskys) in the normal band(norm_thr) normalize the spectrum by inpute ABMag

    Parameters
    ----------
        spectrum: astropy.table, 2 colum, 'WAVELENGTH', 'FLUX'
        norm_thr: astropy.table, 2 colum, 'WAVELENGTH', 'SENSITIVITY'
        sWave: the start of norm_thr
        eWave: the end of norm_thr

    Return:
        the normalization factor   flux of AB system(fix a band  and magnitude ) /the flux of inpute spectrum(fix a band) 
    """
    spectrumi = interpolate.interp1d(spectrum['WAVELENGTH'], spectrum['FLUX'])
    norm_thri = interpolate.interp1d(norm_thr['WAVELENGTH'], norm_thr['SENSITIVITY'])

    x = np.linspace(sWave,eWave, int(eWave)-int(sWave)+1)

    y_spec = spectrumi(x)
    y_thr = norm_thri(x)

    y = y_spec*y_thr

    specAve = np.trapz(y,x)/(eWave-sWave)
    norm = getABMagAverageVal(ABmag=ABMag, norm_thr=norm_thr, sWave=sWave, eWave=eWave)

    if specAve == 0:
        return 0

    # print('specAve = ', specAve)

    return norm / specAve


def tag_sed(h5file, model_tag, teff=5000, logg=2, feh=0):
    model_tag_str = model_tag.decode("utf-8").strip()
    teff_grid = np.unique(h5file["teff"][model_tag_str])
    logg_grid = np.unique(h5file["logg"][model_tag_str])
    feh_grid = np.unique(h5file["feh"][model_tag_str])
    close_teff = teff_grid[np.argmin(np.abs(teff_grid - teff))]
    close_logg = logg_grid[np.argmin(np.abs(logg_grid - logg))]
    if model_tag_str == "koester" or model_tag_str == "MC":
        close_feh = 99
    else:
        close_feh = feh_grid[np.argmin(np.abs(feh_grid - feh))]
    path = model_tag_str + f"_teff_{close_teff:.1f}_logg_{close_logg:.2f}_feh_{close_feh:.1f}"
    wave = np.array(h5file["wave"][model_tag_str][()]).ravel()
    flux = np.array(h5file["sed"][path][()]).ravel()
    return path, wave, flux

from astropy.modeling.models import Gaussian2D
from scipy import signal
def convolveGaussXorders(img=None, sigma = 1):

    offset = int(np.ceil(sigma*10))
    g_size = 2*offset+1

    m_cen = int(g_size/2)

    g_PSF_ = Gaussian2D(1, m_cen, m_cen, sigma, sigma)
    yp, xp = np.mgrid[0:g_size, 0:g_size]
    g_PSF = g_PSF_(xp, yp)
    psf = g_PSF/g_PSF.sum()
    convImg = signal.fftconvolve(img, psf, mode='full', axes=None)
    return convImg, offset