import galsim
import os, sys
import numpy as np
import astropy.constants as cons
from astropy.table import Table
from scipy import interpolate

from ObservationSim.MockObject.MockObject import MockObject
from ObservationSim.MockObject._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG

class Quasar(MockObject):
    def __init__(self, param):
        super().__init__(param)

    def load_SED(self, survey_type, sed_path=None, cosids=None, objtypes=None, sed_templates=None, normFilter=None, target_filt=None):
        if survey_type == "photometric":
            norm_thr_rang_ids = normFilter['SENSITIVITY'] > 0.001
            if sed_templates is None:
                # Read SED data directly
                itype = objtypes[cosids==self.sed_type][0]
                sed_file = os.path.join(sed_path, itype + "_ID%s.sed"%(self.sed_type))
                if not os.path.exists(sed_file):
                    raise ValueError("!!! No SED found.")
                sed_data = Table.read(sed_file, format="ascii")
                wave, flux = sed_data["observedLambda"].data, sed_data["observedFlux"].data
            else:
                # Load SED from templates
                sed_data = sed_templates[self.sed_type]
                # redshift, intrinsic extinction
                sed_data = getObservedSED(
                    sedCat=sed_data, 
                    redshift=self.z, 
                    av=self.param['av'], 
                    redden=self.param['redden'])
                wave, flux = sed_data[0], sed_data[1]
                
            flux_photon = flux * (wave / (cons.h.value * cons.c.value)) * 1e-13
            sed_photon = Table(np.array([wave, flux_photon]).T, names=('WAVELENGTH', 'FLUX'))
            # Get scaling factor for SED
            sedNormFactor = getNormFactorForSpecWithABMAG(ABMag=self.param['mag_use_normal'],
                spectrum=sed_photon,
                norm_thr=normFilter,
                sWave=np.floor(normFilter[norm_thr_rang_ids][0][0]),
                eWave=np.ceil(normFilter[norm_thr_rang_ids][-1][0]))
            sed_photon = np.array([sed_photon['WAVELENGTH'], sed_photon['FLUX']*sedNormFactor]).T
            # Convert to galsim.SED object
            spec = galsim.LookupTable(x=np.array(sed_photon[:, 0]), f=np.array(sed_photon[:, 1]), interpolant='nearest')
            self.sed = galsim.SED(spec, wave_type='A', flux_type='1', fast=False)
            # Get magnitude
            interFlux = integrate_sed_bandpass(sed=self.sed, bandpass=target_filt.bandpass_full)
            self.param['mag_%s'%target_filt.filter_type] = getABMAG(
                interFlux=interFlux, 
                bandpass=target_filt.bandpass_full)
            # print('mag_use_normal = ', self.param['mag_use_normal'])
            # print('mag_%s = '%target_filt.filter_type, self.param['mag_%s'%target_filt.filter_type])

        elif survey_type == "spectroscopic":
            if sed_templates is None:
                self.sedPhotons(sed_path=sed_path, cosids=cosids, objtypes=objtypes)
            else:
                sed_data = sed_templates[self.sed_type]
                sed_data = getObservedSED(
                    sedCat=sed_data, 
                    redshift=self.z, 
                    av=self.param['av'], 
                    redden=self.param['redden'])
                speci = interpolate.interp1d(sed_data[0], sed_data[1])
                lamb = np.arange(2500, 10001 + 0.5, 0.5)
                y = speci(lamb)
                # erg/s/cm2/A --> photo/s/m2/A
                all_sed = y * lamb / (cons.h.value * cons.c.value) * 1e-13
                self.sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))


    def unload_SED(self):
        """(Test) free up SED memory
        """
        del self.sed

    def getGSObj(self, psf, g1=0, g2=0, flux=None, filt=None, tel=None, exptime=150.):
        if flux == None:
            flux = self.getElectronFluxFilt(filt, tel, exptime)
        qso = galsim.Gaussian(sigma=1.e-8, flux=1.)
        qso = qso.withFlux(flux)
        final = galsim.Convolve(psf, qso)
        return final