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

import astropy.io.fits as fitsio

import galsim
import gc
from ObservationSim.MockObject.MockObject import MockObject

from ObservationSim.MockObject._util import magToFlux,VC_A
from ObservationSim.MockObject._util import eObs, integrate_sed_bandpass, getNormFactorForSpecWithABMAG, getObservedSED, getABMAG,convolveGaussXorders

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

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

    def drawObj_multiband(self, tel, pos_img, psf_model, bandpass_list, filt, chip, nphotons_tot=None, g1=0, g2=0, exptime=150., fd_shear=None):
        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)
            self.logger.error(e)
            return False

        nphotons_sum = 0
        photons_list = []
        xmax, ymax = 0, 0

        if self.getMagFilter(filt) <= 15:
            folding_threshold = 5.e-4
        else:
            folding_threshold = 5.e-3
        gsp = galsim.GSParams(folding_threshold=folding_threshold)

        self.real_pos = self.getRealPos(chip.img, global_x=self.posImg.x, global_y=self.posImg.y,
                                        img_real_wcs=self.real_wcs)

        x, y = self.real_pos.x + 0.5, self.real_pos.y + 0.5
        x_nominal = int(np.floor(x + 0.5))
        y_nominal = int(np.floor(y + 0.5))
        dx = x - x_nominal
        dy = y - y_nominal
        offset = galsim.PositionD(dx, dy)

        real_wcs_local = self.real_wcs.local(self.real_pos)

        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)
                self.logger.error(e)
                # return False
                continue

            ratio = sub/full
            if not (ratio == -1 or (ratio != ratio)):
                nphotons = ratio * nphotons_tot
            else:
                # return False
                continue
            nphotons_sum += nphotons

            psf, pos_shear = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, folding_threshold=folding_threshold)

            _gal  = self.param['image']
            galIm = galsim.ImageF(_gal, scale=self.param['pixScale'])
            gal   = galsim.InterpolatedImage(galIm)
            gal   = gal.withFlux(nphotons)
            #gal_shear = galsim.Shear(g1=g1, g2=g2)
            #gal = gal.shear(gal_shear)

            gal = galsim.Convolve(psf, gal)
            if fd_shear is not None:
                gal = gal.shear(fd_shear)

            stamp = gal.drawImage(wcs=real_wcs_local, method='phot', offset=self.offset, save_photons=True)

            xmax = max(xmax, stamp.xmax - stamp.xmin)
            ymax = max(ymax, stamp.ymax - stamp.ymin)

            photons = stamp.photons
            photons.x += x_nominal
            photons.y += y_nominal
            photons_list.append(photons)
            del gal

        # print('xmax = %d, ymax = %d '%(xmax, ymax))

        stamp = galsim.ImageF(int(xmax*1.1), int(ymax*1.1))
        stamp.wcs = real_wcs_local
        stamp.setCenter(x_nominal, y_nominal)
        bounds = stamp.bounds & galsim.BoundsI(0, chip.npix_x - 1, 0, chip.npix_y - 1)

        if bounds.area() > 0:
            chip.img.setOrigin(0, 0)
            stamp[bounds] = chip.img[bounds]
            for i in range(len(photons_list)):
                if i == 0:
                    chip.sensor.accumulate(photons_list[i], stamp)
                else:
                    chip.sensor.accumulate(photons_list[i], stamp, resume=True)

            chip.img[bounds] = stamp[bounds]

            chip.img.setOrigin(chip.bound.xmin, chip.bound.ymin)

        del photons_list
        del stamp
        gc.collect()
        return True, pos_shear