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._util import integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG, tag_sed
from ObservationSim.MockObject.MockObject import MockObject

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