#!/usr/bin/env python
# -*- coding:utf-8 -*-

# @Author: Shuai Feng (hebtu.edu.cn)
# @Time:   2022-09-25

import numpy as np
from scipy.interpolate import interp1d
from astropy.io import fits
from astropy.io import ascii
from astropy import units as u

# ----------------
# Magnitude Module

from scipy.interpolate import interp1d


def Calzetti_Law(wave, Rv=4.05):
    """Dust Extinction Curve by Calzetti et al. (2000)

    Args:
        wave (float): Wavelength
        Rv (float, optional): Extinction curve. Defaults to 4.05.

    Returns:
        float: Extinction value E(B-V)
    """

    wave_number = 1./(wave * 1e-4)
    reddening_curve = np.zeros(len(wave))

    idx = np.logical_and(wave >= 1200, wave <= 6300)
    reddening_curve[idx] = 2.659 * (-2.156 + 1.509 * wave_number[idx] - 0.198 *
                                    (wave_number[idx] ** 2)) + 0.011 * (wave_number[idx] ** 3) + Rv

    idx = np.logical_and(wave >= 6300, wave <= 22000)
    reddening_curve[idx] = 2.659 * (-1.857 + 1.040 * wave_number[idx]) + Rv
    return reddening_curve


def reddening(wave, flux, ebv=0.0, law='calzetti', Rv=4.05):
    """
    Reddening an input spectra through a given reddening curve.

    Args:
        wave (float): Wavelength of input spectra
        flux (float): Flux of input spectra
        ebv (float, optional): Extinction value. Defaults to 0.
        law (str, optional): Extinction curve. Defaults to 'calzetti'.
        Rv (float, optional): _description_. Defaults to 4.05.

    Returns:
        float: Flux of spectra after reddening.
    """
    if law == 'calzetti':
        curve = Calzetti_Law(wave, Rv=Rv)
        fluxNew = flux / (10. ** (0.4 * ebv * curve))
    return fluxNew


def flux_to_mag(wave, flux, path, band='GAIA_bp'):
    """Convert flux of given spectra to magnitude

    Args:
        wave (float): Wavelength
        flux (float): Flux of spectra
        band (str, optional): Filter band name. Defaults to 'GAIA_bp'.

    Returns:
        float: value of magnitude
    """
    # /home/yan/MCI_sim/MCI_input/SED_Code/data
    #

    # parent = os.path.dirname(os.path.realpath(__file__))

    band = ascii.read(path+'MCI_inputData/SED_Code/seddata/' + band + '.dat')
    wave0 = band['col1']
    curv0 = band['col2']

    # Setting the response
    func = interp1d(wave0, curv0)
    response = np.copy(wave)
    ind_extra = (wave > max(wave0)) | (wave < min(wave0))
    response[ind_extra] = 0
    ind_inside = (wave < max(wave0)) & (wave > min(wave0))
    response[ind_inside] = func(wave[ind_inside])

    # Total Flux
    Tflux = np.trapz(flux * response, wave) / np.trapz(response, wave)

    return -2.5 * np.log10(Tflux)


def calibrate(wave, flux, mag, path, band='GAIA_bp'):
    """
    Calibrate the spectra according to the magnitude.

    Args:
        wave (float): Wavelength
        flux (float): Flux of spectra
        mag (float): Input magnitude.
        band (str, optional): Filter band name. Defaults to 'GAIA_bp'.

    Returns:
        float: Flux of calibrated spectra. Units: 1e-17 erg/s/A/cm^2
    """
    inst_mag = flux_to_mag(wave, flux, path, band=band)
    instflux = 10 ** (-0.4 * inst_mag)
    realflux = (mag * u.STmag).to(u.erg/u.s/u.cm**2/u.AA).value

    # Normalization
    flux_ratio = realflux / instflux
    # Units: 10^-17 erg/s/A/cm^2
    flux_calibrate = flux * flux_ratio * 1e17

    return flux_calibrate

# ------------
# SED Template


class Gal_Temp():
    """
    Template of Galaxy SED
    """

    def __init__(self, path):

        #

        # parent = os.path.dirname(os.path.realpath(__file__))
        self.path = path
        hdulist = fits.open(
            self.path+'MCI_inputData/SED_Code/seddata/galaxy_temp.fits')
        self.wave = hdulist[1].data['wave']
        self.flux = hdulist[2].data
        self.age_grid = hdulist[3].data['logAge']
        self.feh_grid = hdulist[3].data['FeH']

    def toMag(self, redshift=0):
        """Calculating magnitude

        Args:
            redshift (float, optional): redshift of spectra. Defaults to 0.
        """
        wave = self.wave * (1 + redshift)
        self.umag = flux_to_mag(wave, self.flux, self.path, band='SDSS_u')
        self.gmag = flux_to_mag(wave, self.flux, self.path, band='SDSS_g')
        self.rmag = flux_to_mag(wave, self.flux, self.path, band='SDSS_r')
        self.imag = flux_to_mag(wave, self.flux, self.path, band='SDSS_i')
        self.zmag = flux_to_mag(wave, self.flux, self.path, band='SDSS_z')


class Star_Temp():
    """
    Template of Stellar SED
    """

    def __init__(self, path):

        #
        self.path = path
        # parent = os.path.dirname(os.path.realpath(__file__))
        # print("获取其父目录——" + parent)  # 从当前文件路径中获取目录
        hdulist = fits.open(
            os.path.join(path, 'MCI_inputData/SED_Code/seddata/stellar_temp.fits'))

        if hdulist[1].data['wave'] == [] or hdulist[2].data == [] or hdulist[3].data['Teff'] == [] or hdulist[3].data['FeH'] == []:
            self.log.error(
                'Data error in tellar_temp.fits file, please check!!!')
            raise ValueError(
                'Data error in tellar_temp.fits file, please check!!!')

        self.wave = hdulist[1].data['wave']
        self.flux = hdulist[2].data
        self.Teff_grid = hdulist[3].data['Teff']
        self.FeH_grid = hdulist[3].data['FeH']

        self.bpmag = flux_to_mag(self.wave, self.flux, path, band='GAIA_bp')
        self.rpmag = flux_to_mag(self.wave, self.flux, path, band='GAIA_rp')

    def toMag(self):
        wave = self.wave
        self.bpmag = flux_to_mag(wave, self.flux, self.path, band='GAIA_bp')
        self.rpmag = flux_to_mag(wave, self.flux, self.path, band='GAIA_rp')

# -------------
# SED Modelling


def Model_Stellar_SED(wave, bp, rp, temp):
    """Modelling stellar SED based on bp, rp magnitude

    Args:
        wave (float): Wavelength
        bp (float): Magnitude of GAIA BP band
        rp (float): Magnitude of GAIA RP band
        temp (class): Class of stellar template

    Returns:
        float array: Spectral energy distribution of stellar SED,
        which have the same length to the input wave
        """

    color0 = bp - rp
    colors = temp.bpmag - temp.rpmag

    idx = np.argmin(np.abs(colors - color0))
    flux0 = temp.flux[idx]
    flux1 = np.interp(wave, temp.wave, flux0)
    flux = calibrate(wave, flux1, rp, band='GAIA_rp')
    return flux


def Model_Galaxy_SED(wave, ugriz, z, temp, path):
    """Modelling galaxy SED based on u,g,r,i,z magnitude

    Args:
        wave (float): Wavelength
        ugriz (float, array): The array of magnitude of SDSS ugriz band
        z (float): Redshift
        temp (class): Class of gaalxy template

    Returns:
        float array: Spectral energy distribution of stellar SED,
        which have the same length to the input wave
    """
    sed = 10. ** (-0.4 * ugriz)
    sed = sed / sed[2]
    ntemp = len(temp.rmag)
    dmag = np.zeros(ntemp)
    for j in range(ntemp):
        ugriz0 = np.array([temp.umag[j], temp.gmag[j], temp.rmag[j],
                           temp.imag[j], temp.zmag[j]])
        sed0 = 10. ** (-0.4 * ugriz0)
        sed0 = sed0 / sed0[2]
        dmag[j] = np.sum(np.abs(sed - sed0))

    idx = np.argmin(dmag)
    flux0 = temp.flux[idx]

    # Effect of E(B-V)
    ri0 = ugriz[2] - ugriz[3]
    ri = temp.rmag - temp.imag
    dri = ri0 - ri[idx]
    Alambda = Calzetti_Law(np.array([6213 / (1 + z), 7625 / (1 + z)]))
    eri0 = (Alambda[0] - Alambda[1])
    ebv = dri/eri0
    if ebv < 0:
        ebv = 0
    if ebv > 0.5:
        ebv = 0.5
    flux1 = reddening(temp.wave, flux0, ebv=ebv)

    flux2 = np.interp(wave, temp.wave * (1 + z), flux1)
    flux = calibrate(wave, flux2, ugriz[2], path, band='SDSS_r')
    return flux
