import galsim
import os, sys
import numpy as np
import astropy.constants as cons
from astropy.table import Table
from ._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG, tag_sed
from .SpecDisperser import SpecDisperser
from .MockObject import MockObject
from scipy import interpolate

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

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

    def load_SED(self, survey_type, normFilter=None, target_filt=None, sed_lib=None, sed_path=None):
        if survey_type == "photometric":
            norm_thr_rang_ids = normFilter['SENSITIVITY'] > 0.001
            # spec_lambda = Table.read(sed_path, path=f"/SED/wave_{self.model_tag}")
            # spec_flux = Table.read(sed_path, path=f"/SED/{self.sed_type}")
            # wave, flux = spec_lambda['col0'].data, spec_flux['col0'].data
            _, wave, flux = tag_sed(
                h5file=sed_lib, 
                model_tag=self.param['model_tag'], 
                teff=self.param['teff'],
                logg=self.param['logg'],
                feh=self.param['feh'])
            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":
            # self.sedPhotons(sed_path=sed_path)
            self.sedPhotons(sed_lib=sed_lib)

    def sedPhotons(self, sed_path=None, sed_lib=None):
        # spec_lambda = Table.read(sed_path, path=f"/SED/wave_{self.model_tag}")
        # spec_flux = Table.read(sed_path, path=f"/SED/{self.sed_type}")
        _, w_orig, f_orig = tag_sed(
                h5file=sed_lib, 
                model_tag=self.param['model_tag'], 
                teff=self.param['teff'],
                logg=self.param['logg'],
                feh=self.param['feh'])
        # spec_data = {}
        # f_orig = spec_flux["col0"].data
        # w_orig = spec_lambda["col0"].data
        speci = interpolate.interp1d(w_orig, f_orig)
        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 getGSObj(self, psf, g1=0, g2=0, flux=None, filt=None, tel=None, exptime=150.):
        if flux == None:
            flux = self.getElectronFluxFilt(filt, tel, exptime)
        # star = galsim.Gaussian(sigma=1.e-8, flux=1.)
        star = galsim.DeltaFunction()
        star = star.withFlux(flux)
        final = galsim.Convolve(psf, star)
        return final
        
    def getGSObj_multiband(self, tel, psf_list, bandpass_list, filt, nphotons_tot=None, g1=0, g2=0, exptime=150.):
        if len(psf_list) != len(bandpass_list):
            raise ValueError("!!!The number of PSF profiles and the number of bandpasses must be equal.")
        objs = []
        if nphotons_tot == None:
            nphotons_tot = self.getElectronFluxFilt(filt, tel, exptime)

        try:
            full = integrate_sed_bandpass(sed=self.sed, bandpass=filt.bandpass_full)
        except Exception as e:
            print(e)
            return -1

        for i in range(len(bandpass_list)):
            bandpass = bandpass_list[i]
            try:
                sub = integrate_sed_bandpass(sed=self.sed, bandpass=bandpass)
            except Exception as e:
                print(e)
                return -1
        
            ratio = sub/full

            if not (ratio == -1 or (ratio != ratio)):
                nphotons = ratio * nphotons_tot
            else:
                return -1
            star = galsim.DeltaFunction()
            star = star.withFlux(nphotons)
            star = galsim.Convolve(psf, star)
            objs.append(star)
        final = galsim.Sum(objs)
        return final