PSFInterpSLS.py

'''
PSF interpolation for CSST-Sim

NOTE: [iccd, iwave, ipsf] are counted from 1 to n, but [tccd, twave, tpsf] are counted from 0 to n-1
'''

import yaml
import sys
import time
import copy
import numpy as np
import scipy.spatial as spatial
import galsim
import h5py

from observation_sim.instruments import Filter, FilterParam, Chip
from observation_sim.psf.PSFModel import PSFModel
from observation_sim.instruments.chip import chip_utils
import os
from astropy.io import fits

from astropy.modeling.models import Gaussian2D
from scipy import signal, interpolate
import datetime
import gc
# from jax import numpy as jnp

LOG_DEBUG = False  # ***#
NPSF = 900  # ***# 30*30
PIX_SIZE_MICRON = 5.  # ***# in microns


### find neighbors-KDtree###
# def findNeighbors(tx, ty, px, py, dr=0.1, dn=1, OnlyDistance=True):
#     """
#     find nearest neighbors by 2D-KDTree
#
#     Parameters:
#         tx, ty (float, float): a given position
#         px, py (numpy.array, numpy.array): position data for tree
#         dr (float-optional): distance
#         dn (int-optional): nearest-N
#         OnlyDistance (bool-optional): only use distance to find neighbors. Default: True
#
#     Returns:
#         dataq (numpy.array): index
#     """
#     datax = px
#     datay = py
#     tree = spatial.KDTree(list(zip(datax.ravel(), datay.ravel())))
#
#     dataq=[]
#     rr = dr
#     if OnlyDistance == True:
#         dataq = tree.query_ball_point([tx, ty], rr)
#     if OnlyDistance == False:
#         while len(dataq) < dn:
#             dataq = tree.query_ball_point([tx, ty], rr)
#             rr += dr
#         dd = np.hypot(datax[dataq]-tx, datay[dataq]-ty)
#         ddSortindx = np.argsort(dd)
#         dataq = np.array(dataq)[ddSortindx[0:dn]]
#     return dataq
#
# ###find neighbors-hoclist###
# def hocBuild(partx, party, nhocx, nhocy, dhocx, dhocy):
#     if np.max(partx) > nhocx*dhocx:
#         print('ERROR')
#         sys.exit()
#     if np.max(party) > nhocy*dhocy:
#         print('ERROR')
#         sys.exit()
#
#     npart  = partx.size
#     hoclist= np.zeros(npart, dtype=np.int32)-1
#     hoc = np.zeros([nhocy, nhocx], dtype=np.int32)-1
#     for ipart in range(npart):
#         ix = int(partx[ipart]/dhocx)
#         iy = int(party[ipart]/dhocy)
#         hoclist[ipart] = hoc[iy, ix]
#         hoc[iy, ix] = ipart
#     return hoc, hoclist
#
# def hocFind(px, py, dhocx, dhocy, hoc, hoclist):
#     ix = int(px/dhocx)
#     iy = int(py/dhocy)
#
#     neigh=[]
#     it = hoc[iy, ix]
#     while it != -1:
#         neigh.append(it)
#         it = hoclist[it]
#     return neigh
#
# def findNeighbors_hoclist(px, py, tx=None,ty=None, dn=4, hoc=None, hoclist=None):
#     nhocy = nhocx = 20
#
#     pxMin = np.min(px)
#     pxMax = np.max(px)
#     pyMin = np.min(py)
#     pyMax = np.max(py)
#
#     dhocx = (pxMax - pxMin)/(nhocx-1)
#     dhocy = (pyMax - pyMin)/(nhocy-1)
#     partx = px - pxMin +dhocx/2
#     party = py - pyMin +dhocy/2
#
#     if hoc is None:
#         hoc, hoclist = hocBuild(partx, party, nhocx, nhocy, dhocx, dhocy)
#         return hoc, hoclist
#
#     if hoc is not None:
#         tx = tx - pxMin +dhocx/2
#         ty = ty - pyMin +dhocy/2
#         itx = int(tx/dhocx)
#         ity = int(ty/dhocy)
#
#         ps = [-1, 0, 1]
#         neigh=[]
#         for ii in range(3):
#             for jj in range(3):
#                 ix = itx + ps[ii]
#                 iy = ity + ps[jj]
#                 if ix < 0:
#                     continue
#                 if iy < 0:
#                     continue
#                 if ix > nhocx-1:
#                     continue
#                 if iy > nhocy-1:
#                     continue
#
#                 #neightt = myUtil.hocFind(ppx, ppy, dhocx, dhocy, hoc, hoclist)
#                 it = hoc[iy, ix]
#                 while it != -1:
#                     neigh.append(it)
#                     it = hoclist[it]
#                 #neigh.append(neightt)
#         #ll = [i for k in neigh for i in k]
#         if dn != -1:
#             ptx = np.array(partx[neigh])
#             pty = np.array(party[neigh])
#             dd  = np.hypot(ptx-tx, pty-ty)
#             idx = np.argsort(dd)
#             neigh= np.array(neigh)[idx[0:dn]]
#         return neigh
#
#
# ###PSF-IDW###
# def psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=None, hoclist=None, PSFCentroidWgt=False):
#     """
#     psf interpolation by IDW
#
#     Parameters:
#         px, py (float, float): position of the target
#         PSFMat (numpy.array): image
#         cen_col, cen_row (numpy.array, numpy.array): potions of the psf centers
#         IDWindex (int-optional): the power index of IDW
#         OnlyNeighbors (bool-optional): only neighbors are used for psf interpolation
#
#     Returns:
#         psfMaker (numpy.array)
#     """
#
#     minimum_psf_weight = 1e-8
#     ref_col = px
#     ref_row = py
#
#     ngy, ngx = PSFMat[0, :, :].shape
#     npsf = PSFMat[:, :, :].shape[0]
#     psfWeight = np.zeros([npsf])
#
#     if OnlyNeighbors == True:
#         if hoc is None:
#             neigh = findNeighbors(px, py, cen_col, cen_row, dr=5., dn=4, OnlyDistance=False)
#         if hoc is not None:
#             neigh = findNeighbors_hoclist(cen_col, cen_row, tx=px,ty=py, dn=4, hoc=hoc, hoclist=hoclist)
#
#         neighFlag = np.zeros(npsf)
#         neighFlag[neigh] = 1
#
#     for ipsf in range(npsf):
#         if OnlyNeighbors == True:
#             if neighFlag[ipsf] != 1:
#                 continue
#
#         dist = np.sqrt((ref_col - cen_col[ipsf])**2 + (ref_row - cen_row[ipsf])**2)
#         if IDWindex == 1:
#             psfWeight[ipsf] = dist
#         if IDWindex == 2:
#             psfWeight[ipsf] = dist**2
#         if IDWindex == 3:
#             psfWeight[ipsf] = dist**3
#         if IDWindex == 4:
#             psfWeight[ipsf] = dist**4
#         psfWeight[ipsf] = max(psfWeight[ipsf], minimum_psf_weight)
#         psfWeight[ipsf] = 1./psfWeight[ipsf]
#     psfWeight /= np.sum(psfWeight)
#
#     psfMaker  = np.zeros([ngy, ngx], dtype=np.float32)
#     for ipsf in range(npsf):
#         if OnlyNeighbors == True:
#             if neighFlag[ipsf] != 1:
#                 continue
#
#         iPSFMat = PSFMat[ipsf, :, :].copy()
#         ipsfWeight = psfWeight[ipsf]
#
#         psfMaker += iPSFMat * ipsfWeight
#     psfMaker /= np.nansum(psfMaker)
#
#     return psfMaker


### define PSFInterp###
class PSFInterpSLS(PSFModel):
    def __init__(self, chip, filt, PSF_data_prefix="", sigSpin=0, psfRa=0.15, pix_size=0.005):
        if LOG_DEBUG:
            print('===================================================')
            print('DEBUG: psf module for csstSim '
                  + time.strftime("(%Y-%m-%d %H:%M:%S)", time.localtime()), flush=True)
            print('===================================================')

        self.sigSpin = sigSpin
        self.sigGauss = psfRa
        self.grating_ids = chip_utils.getChipSLSGratingID(chip.chipID)
        _, self.grating_type = chip.getChipFilter(chipID=chip.chipID)
        self.data_folder = PSF_data_prefix
        self.getPSFDataFromFile(filt)
        self.pixsize = pix_size  # um

    def getPSFDataFromFile(self, filt):
        gratingInwavelist = {'GU': 0, 'GV': 1, 'GI': 2}
        grating_orders = ['0', '1']
        waveListFn = self.data_folder + '/wavelist.dat'
        wavelists = np.loadtxt(waveListFn)
        self.waveList = wavelists[:, gratingInwavelist[self.grating_type]]
        bandranges = np.zeros([4, 2])
        midBand = (self.waveList[0:3] + self.waveList[1:4])/2.*10000.
        bandranges[0, 0] = filt.blue_limit
        bandranges[1:4, 0] = midBand
        bandranges[0:3, 1] = midBand
        bandranges[3, 1] = filt.red_limit

        self.bandranges = bandranges

        self.grating1_data = {}
        g_folder = self.data_folder + '/' + self.grating_ids[0] + '/'
        for g_order in grating_orders:
            g_folder_order = g_folder + 'PSF_Order_' + g_order + '/'
            grating_order_data = {}
            for bandi in [1, 2, 3, 4]:
                subBand_data = {}
                subBand_data['bandrange'] = bandranges[bandi-1]
                final_folder = g_folder_order + str(bandi) + '/'
                print(final_folder)
                pca_fs = os.listdir(final_folder)
                for fname in pca_fs:
                    if ('_PCs.fits' in fname) and (fname[0] != '.'):
                        fname_ = final_folder + fname
                        hdu = fits.open(fname_)
                        subBand_data['band_data'] = hdu
                grating_order_data['band'+str(bandi)] = subBand_data
            self.grating1_data['order'+g_order] = grating_order_data

        self.grating2_data = {}
        g_folder = self.data_folder + '/' + self.grating_ids[1] + '/'
        for g_order in grating_orders:
            g_folder_order = g_folder + 'PSF_Order_' + g_order + '/'
            grating_order_data = {}
            for bandi in [1, 2, 3, 4]:
                subBand_data = {}
                subBand_data['bandrange'] = bandranges[bandi - 1]
                final_folder = g_folder_order + str(bandi) + '/'
                print(final_folder)
                pca_fs = os.listdir(final_folder)
                for fname in pca_fs:
                    if ('_PCs.fits' in fname) and (fname[0] != '.'):
                        fname_ = final_folder + fname
                        hdu = fits.open(fname_)
                        subBand_data['band_data'] = hdu
                grating_order_data['band' + str(bandi)] = subBand_data
            self.grating2_data['order' + g_order] = grating_order_data

    #
    #
    #
    # def _getPSFwave(self, iccd, PSF_data_file, PSF_data_prefix):
    #     # fq = h5py.File(PSF_data_file+'/' +PSF_data_prefix +'psfCube_ccd{:}.h5'.format(iccd), 'r')
    #     fq = h5py.File(PSF_data_file+'/' +PSF_data_prefix +'psfCube_{:}.h5'.format(iccd), 'r')
    #     nwave = len(fq.keys())
    #     fq.close()
    #     return nwave
    #
    #
    # def _loadPSF(self, iccd, PSF_data_file, PSF_data_prefix):
    #     psfSet = []
    #     # fq = h5py.File(PSF_data_file+'/' +PSF_data_prefix +'psfCube_ccd{:}.h5'.format(iccd), 'r')
    #     fq = h5py.File(PSF_data_file+'/' +PSF_data_prefix +'psfCube_{:}.h5'.format(iccd), 'r')
    #     for ii in range(self.nwave):
    #         iwave = ii+1
    #         psfWave = []
    #
    #         fq_iwave = fq['w_{:}'.format(iwave)]
    #         for jj in range(self.npsf):
    #             ipsf = jj+1
    #             psfInfo = {}
    #             psfInfo['wavelength']= fq_iwave['wavelength'][()]
    #
    #             fq_iwave_ipsf = fq_iwave['psf_{:}'.format(ipsf)]
    #             psfInfo['pixsize']   = PIX_SIZE_MICRON
    #             psfInfo['field_x']   = fq_iwave_ipsf['field_x'][()]
    #             psfInfo['field_y']   = fq_iwave_ipsf['field_y'][()]
    #             psfInfo['image_x']   = fq_iwave_ipsf['image_x'][()]
    #             psfInfo['image_y']   = fq_iwave_ipsf['image_y'][()]
    #             psfInfo['centroid_x']= fq_iwave_ipsf['cx'][()]
    #             psfInfo['centroid_y']= fq_iwave_ipsf['cy'][()]
    #             psfInfo['psfMat']    = fq_iwave_ipsf['psfMat'][()]
    #
    #             psfWave.append(psfInfo)
    #         psfSet.append(psfWave)
    #     fq.close()
    #
    #     if LOG_DEBUG:
    #         print('psfSet has been loaded:', flush=True)
    #         print('psfSet[iwave][ipsf][keys]:', psfSet[0][0].keys(), flush=True)
    #     return psfSet
    #
    #
    # def _findWave(self, bandpass):
    #     if isinstance(bandpass,int):
    #         twave = bandpass
    #         return twave
    #
    #     for twave in range(self.nwave):
    #         bandwave = self.PSF_data[twave][0]['wavelength']
    #         if bandpass.blue_limit < bandwave and bandwave < bandpass.red_limit:
    #             return twave
    #     return -1
    #
    #

    def convolveWithGauss(self, img=None, sigma=1):

        offset = int(np.ceil(sigma * 3))
        g_size = 2 * offset + 1
        m_cen = int(g_size / 2)
        print('-----', g_size)
        g_PSF_ = Gaussian2D(1, m_cen, m_cen, sigma, sigma)
        yp, xp = np.mgrid[0:g_size, 0:g_size]
        g_PSF = g_PSF_(xp, yp)
        psf = g_PSF / g_PSF.sum()
        convImg = signal.fftconvolve(img, psf, mode='full', axes=None)
        convImg = convImg/np.sum(convImg)
        return convImg

    def get_PSF(self, chip, pos_img_local=[1000, 1000], bandNo=1, galsimGSObject=True, folding_threshold=5.e-3, g_order='A', grating_split_pos=3685):
        """
        Get the PSF at a given image position

        Parameters:
            chip: A 'Chip' object representing the chip we want to extract PSF from.
            pos_img: A 'galsim.Position' object representing the image position.
            bandpass: A 'galsim.Bandpass' object representing the wavelength range.
            pixSize: The pixels size of psf matrix
            findNeighMode: 'treeFind' or 'hoclistFind'
        Returns:
            PSF: A 'galsim.GSObject'.
        """
        order_IDs = {'A': '1', 'B': '0', 'C': '0', 'D': '0', 'E': '0'}
        contam_order_sigma = {'C': 0.28032344707964174,
                              'D': 0.39900182912061344, 'E': 1.1988309797685412}  # arcsec
        x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
        y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.
        # print(pos_img.x - x_start)
        pos_img_x = pos_img_local[0] + x_start
        pos_img_y = pos_img_local[1] + y_start
        pos_img = galsim.PositionD(pos_img_x, pos_img_y)
        if pos_img_local[0] < grating_split_pos:
            psf_data = self.grating1_data
        else:
            psf_data = self.grating2_data

        grating_order = order_IDs[g_order]
        # if grating_order in ['-2','-1','2']:
        #     grating_order = '1'

        # if grating_order in ['0', '1']:
        psf_order = psf_data['order'+grating_order]
        psf_order_b = psf_order['band'+str(bandNo)]
        psf_b_dat = psf_order_b['band_data']
        pos_p = psf_b_dat[1].data
        pc_coeff = psf_b_dat[2].data
        pcs = psf_b_dat[0].data
        # print(max(pos_p[:,0]), min(pos_p[:,0]),max(pos_p[:,1]), min(pos_p[:,1]))
        # print(chip.x_cen, chip.y_cen)
        # print(pos_p)
        px = pos_img.x*chip.pix_size
        py = pos_img.y*chip.pix_size

        dist2 = (pos_p[:, 1] - px)*(pos_p[:, 1] - px) + \
            (pos_p[:, 0] - py)*(pos_p[:, 0] - py)
        temp_sort_dist = np.zeros([dist2.shape[0], 2])
        temp_sort_dist[:, 0] = np.arange(0, dist2.shape[0], 1)
        temp_sort_dist[:, 1] = dist2
        # print(temp_sort_dist)
        dits2_sortlist = sorted(temp_sort_dist, key=lambda x: x[1])
        # print(dits2_sortlist)
        nearest4p = np.zeros([4, 2])
        pc_coeff_4p = np.zeros([pc_coeff.data.shape[0], 4])

        for i in np.arange(4):
            smaller_ids = int(dits2_sortlist[i][0])
            nearest4p[i, 0] = pos_p[smaller_ids, 1]
            nearest4p[i, 1] = pos_p[smaller_ids, 0]
            pc_coeff_4p[:, i] = pc_coeff[:, smaller_ids]
        idw_dist = 1/(np.sqrt((px-nearest4p[:, 0]) * (px-nearest4p[:, 0]) + (
            py-nearest4p[:, 1]) * (py-nearest4p[:, 1])))

        coeff_int = np.zeros(pc_coeff.data.shape[0])
        for i in np.arange(4):
            coeff_int = coeff_int + pc_coeff_4p[:, i]*idw_dist[i]
        coeff_int = coeff_int / np.sum(coeff_int)

        npc = 10
        m_size = int(pcs.shape[0]**0.5)
        PSF_int = np.dot(pcs[:, 0:npc], coeff_int[0:npc]
                         ).reshape(m_size, m_size)

        # PSF_int = PSF_int/np.sum(PSF_int)
        PSF_int_trans = np.flipud(np.fliplr(PSF_int))
        PSF_int_trans = np.fliplr(PSF_int_trans.T)
        # PSF_int_trans = np.abs(PSF_int_trans)
        # ids_szero = PSF_int_trans<0
        # PSF_int_trans[ids_szero] = 0
        # print(PSF_int_trans[ids_szero].shape[0],PSF_int_trans.shape)
        PSF_int_trans = PSF_int_trans/np.sum(PSF_int_trans)
        ###DEBGU
        ids_szero = PSF_int_trans<0
        n01 = PSF_int_trans[ids_szero].shape[0]

        n1 = np.sum(np.isinf(PSF_int_trans))
        n2 = np.sum(np.isnan(PSF_int_trans))
        if n1>0 or n2>0:
            print("DEBUG: PSFInterpSLS, inf:%d, nan:%d, 0 num:%d"%(n1, n2, n01))

        ####
        # from astropy.io import fits
        # fits.writeto(str(bandNo) + '_' + g_order+ '_psf_o.fits', PSF_int_trans)

        # if g_order in ['C','D','E']:
        #     g_simgma = contam_order_sigma[g_order]/pixel_size_arc
        #     PSF_int_trans = self.convolveWithGauss(PSF_int_trans,g_simgma)
        # n_m_size = int(m_size/2)
        #
        # n_PSF_int = np.zeros([n_m_size, n_m_size])
        #
        # for i in np.arange(n_m_size):
        #     for j in np.arange(n_m_size):
        #         n_PSF_int[i,j] = np.sum(PSF_int[2*i:2*i+2, 2*j:2*j+2])
        #
        # n_PSF_int = n_PSF_int/np.sum(n_PSF_int)

        # chip.img = galsim.ImageF(chip.npix_x, chip.npix_y)
        # chip.img.wcs = galsim.wcs.AffineTransform
        if galsimGSObject:

            # imPSFt = np.zeros([257,257])
            # imPSFt[0:256, 0:256] = imPSF
            # # imPSFt[120:130, 0:256] = 1.

            pixel_size_arc = np.rad2deg(self.pixsize * 1e-3 / 28) * 3600
            img = galsim.ImageF(PSF_int_trans, scale=pixel_size_arc)
            gsp = galsim.GSParams(folding_threshold=folding_threshold)
            # TEST: START
            # Use sheared PSF to test the PSF orientation
            # self.psf = galsim.InterpolatedImage(img, gsparams=gsp).shear(g1=0.8, g2=0.)
            # TEST: END
            self.psf = galsim.InterpolatedImage(img, gsparams=gsp)
            # if g_order in ['C','D','E']:
            #     add_psf  = galsim.Gaussian(sigma=contam_order_sigma[g_order], flux=1.0)
            #     self.psf = galsim.Convolve(self.psf, add_psf)
            wcs = chip.img.wcs.local(pos_img)
            scale = galsim.PixelScale(0.074)
            self.psf = wcs.toWorld(scale.toImage(
                self.psf), image_pos=(pos_img))

            # return self.PSFspin(x=px/0.01, y=py/0.01)
            return self.psf, galsim.Shear(e=0., beta=(np.pi/2)*galsim.radians)

        return PSF_int_trans, PSF_int

    def get_PSF_AND_convolve_withsubImg(self, chip, cutImg=None, pos_img_local=[1000, 1000], bandNo=1, g_order='A', grating_split_pos=3685):
        """
        Get the PSF at a given image position

        Parameters:
            chip: A 'Chip' object representing the chip we want to extract PSF from.
            pos_img: A 'galsim.Position' object representing the image position.
            bandpass: A 'galsim.Bandpass' object representing the wavelength range.
            pixSize: The pixels size of psf matrix
            findNeighMode: 'treeFind' or 'hoclistFind'
        Returns:
            PSF: A 'galsim.GSObject'.
        """
        order_IDs = {'A': '1', 'B': '0', 'C': '0', 'D': '0', 'E': '0'}
        contam_order_sigma = {'C': 0.28032344707964174,
                              'D': 0.39900182912061344, 'E': 1.1988309797685412}  # arcsec
        x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
        y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.
        # print(pos_img.x - x_start)
        # pos_img_x = pos_img_local[0] + x_start
        # pos_img_y = pos_img_local[1] + y_start
        # pos_img = galsim.PositionD(pos_img_x, pos_img_y)
        # centerPos_local = cutImg.ncol/2.
        if pos_img_local[0] < grating_split_pos:
            psf_data = self.grating1_data
        else:
            psf_data = self.grating2_data

        grating_order = order_IDs[g_order]
        # if grating_order in ['-2','-1','2']:
        #     grating_order = '1'

        # if grating_order in ['0', '1']:
        psf_order = psf_data['order'+grating_order]
        psf_order_b = psf_order['band'+str(bandNo)]
        psf_b_dat = psf_order_b['band_data']
        # pos_p = psf_b_dat[1].data
        pos_p = psf_b_dat[1].data/chip.pix_size - np.array([y_start, x_start])

        pc_coeff = psf_b_dat[2].data
        pcs = psf_b_dat[0].data

        npc = 10
        m_size = int(pcs.shape[0]**0.5)
        sumImg = np.sum(cutImg.array)
        tmp_img = cutImg*0
        for j in np.arange(npc):
            X_ = np.hstack((pos_p[:,1].flatten()[:, None], pos_p[:,0].flatten()[:, None]),dtype=np.float32)
            Z_ = (pc_coeff[j].astype(np.float32)).flatten()
            # print(pc_coeff[j].shape[0], pos_p[:,1].shape[0], pos_p[:,0].shape[0])
            cx_len = int(chip.npix_x)
            cy_len = int(chip.npix_y)
            n_x = np.arange(0, cx_len, 1, dtype = int)
            n_y = np.arange(0, cy_len, 1, dtype = int)
            M, N = np.meshgrid(n_x, n_y)
            # t1=datetime.datetime.now()
    #         U = interpolate.griddata(X_, Z_, (M[0:cy_len, 0:cx_len],N[0:cy_len, 0:cx_len]),
    # method='nearest',fill_value=1.0)
            b_img = galsim.Image(cx_len, cy_len)
            b_img.setOrigin(0,0)
            bounds = cutImg.bounds & b_img.bounds
            if bounds.area() == 0:
                continue

            # ys = cutImg.ymin
            # if ys < 0: 
            #     ys = 0
            # ye = cutImg.ymin+cutImg.nrow
            # if ye >= cy_len-1: 
            #     ye = cy_len-1
            # if ye - ys <=0:
            #     continue
            # xs = cutImg.xmin
            # if xs < 0: 
            #     xs = 0
            # xe = cutImg.xmin+cutImg.ncol
            # if xe >= cx_len-1: 
            #     xe = cx_len-1
            # if xe - xs <=0:
            #     continue
            ys = bounds.ymin
            ye = bounds.ymax+1
            xs = bounds.xmin
            xe = bounds.xmax+1
            U = interpolate.griddata(X_, Z_, (M[ys:ye, xs:xe],N[ys:ye, xs:xe]),
    method='nearest',fill_value=1.0)
            # t2=datetime.datetime.now()
            
            # print("time interpolate:", t2-t1)

            # if U.shape != cutImg.array.shape:
            #     print('DEBUG:SHAPE',cutImg.ncol,cutImg.nrow,cutImg.xmin, cutImg.ymin)
            #     continue
            img_tmp = cutImg
            img_tmp[bounds] = img_tmp[bounds]*U
            psf = pcs[:, j].reshape(m_size, m_size)
            tmp_img = tmp_img + signal.fftconvolve(img_tmp.array, psf, mode='same', axes=None)

            # t3=datetime.datetime.now()
            # print("time convole:", t3-t2)
            del U
            del img_tmp
        if np.sum(tmp_img.array)==0:
            tmp_img = cutImg
        else:
            tmp_img = tmp_img/np.sum(tmp_img.array)*sumImg
        return tmp_img


    def convolveFullImgWithPCAPSF(self, chip, folding_threshold=5.e-3):
        keys_L1= chip_utils.getChipSLSGratingID(chip.chipID)
        # keys_L2 = ['order-2','order-1','order0','order1','order2']
        keys_L2 = ['order0','order1']
        keys_L3 = ['w1','w2','w3','w4']

        npca = 10

        x_start = chip.x_cen/chip.pix_size - chip.npix_x / 2.
        y_start = chip.y_cen/chip.pix_size - chip.npix_y / 2.

        for i,gt in enumerate(keys_L1):
            psfCo = self.grating1_data
            if i > 0:
                psfCo = self.grating2_data
            for od in keys_L2:
                psfCo_L2 = psfCo['order1']
                if od in ['order-2','order-1','order0','order2']:
                    psfCo_L2 = psfCo['order0']
                for w in keys_L3:
                    img = chip.img_stack[gt][od][w]
                    pcs = psfCo_L2['band'+w[1]]['band_data'][0].data
                    pos_p = psfCo_L2['band'+w[1]]['band_data'][1].data/chip.pix_size - np.array([y_start, x_start])
                    pc_coeff = psfCo_L2['band'+w[1]]['band_data'][2].data
                    # print("DEBUG-----------",np.max(pos_p[:,1]),np.min(pos_p[:,1]), np.max(pos_p[:,0]),np.min(pos_p[:,0]))
                    sum_img = np.sum(img.array)
                    
                    
                    # coeff_mat = np.zeros([npca, chip.npix_y, chip.npix_x])
                    # for m in np.arange(chip.npix_y):
                    #     for n in np.arange(chip.npix_x):
                    #         px = n
                    #         py = m
                            
                    #         dist2 = (pos_p[:, 1] - px)*(pos_p[:, 1] - px) + (pos_p[:, 0] - py)*(pos_p[:, 0] - py)
                    #         temp_sort_dist = np.zeros([dist2.shape[0], 2])
                    #         temp_sort_dist[:, 0] = np.arange(0, dist2.shape[0], 1)
                    #         temp_sort_dist[:, 1] = dist2
                    #         # print(temp_sort_dist)
                    #         dits2_sortlist = sorted(temp_sort_dist, key=lambda x: x[1])
                    #         # print(dits2_sortlist)
                    #         nearest4p = np.zeros([4, 3])
                    #         pc_coeff_4p = np.zeros([npca, 4])

                    #         for i in np.arange(4):
                    #             smaller_ids = int(dits2_sortlist[i][0])
                    #             nearest4p[i, 0] = pos_p[smaller_ids, 1]
                    #             nearest4p[i, 1] = pos_p[smaller_ids, 0]
                    #             # print(pos_p[smaller_ids, 1],pos_p[smaller_ids, 0])
                    #             nearest4p[i, 2] = dits2_sortlist[i][1]
                    #             pc_coeff_4p[:, i] = pc_coeff[npca, smaller_ids]
                    #         # idw_dist = 1/(np.sqrt((px-nearest4p[:, 0]) * (px-nearest4p[:, 0]) + (
                    #         #     py-nearest4p[:, 1]) * (py-nearest4p[:, 1])))
                    #         idw_dist = 1/(np.sqrt(nearest4p[:, 2]))

                    #         coeff_int = np.zeros(npca)
                    #         for i in np.arange(4):
                    #             coeff_int = coeff_int + pc_coeff_4p[:, i]*idw_dist[i]
                    #         coeff_mat[:, m, n] = coeff_int

                    m_size = int(pcs.shape[0]**0.5)
                    tmp_img = np.zeros_like(img.array,dtype=np.float32)
                    for j in np.arange(npca):
                        print(gt, od, w, j)
                        X_ = np.hstack((pos_p[:,1].flatten()[:, None], pos_p[:,0].flatten()[:, None]),dtype=np.float32)
                        Z_ = (pc_coeff[j].astype(np.float32)).flatten()
                        # print(pc_coeff[j].shape[0], pos_p[:,1].shape[0], pos_p[:,0].shape[0])
                        sub_size = 4
                        cx_len = int(chip.npix_x/sub_size)
                        cy_len = int(chip.npix_y/sub_size)
                        n_x = np.arange(0, chip.npix_x, sub_size, dtype = int)
                        n_y = np.arange(0, chip.npix_y, sub_size, dtype = int)

                        M, N = np.meshgrid(n_x, n_y)
                        t1=datetime.datetime.now()
                #         U = interpolate.griddata(X_, Z_, (M[0:cy_len, 0:cx_len],N[0:cy_len, 0:cx_len]),
                # method='nearest',fill_value=1.0)
                        U1 = interpolate.griddata(X_, Z_, (M, N),
                method='nearest',fill_value=1.0)
                        U = np.zeros_like(chip.img.array, dtype=np.float32)
                        for mi in np.arange(cy_len):
                            for mj in np.arange(cx_len):
                                U[mi*sub_size:(mi+1)*sub_size, mj*sub_size:(mj+1)*sub_size]=U1[mi,mj]
                        t2=datetime.datetime.now()
                        
                        print("time interpolate:", t2-t1)

                        img_tmp = img.array*U
                        psf = pcs[:, j].reshape(m_size, m_size)
                        tmp_img = tmp_img + signal.fftconvolve(img_tmp, psf, mode='same', axes=None)

                        t3=datetime.datetime.now()
                        print("time convole:", t3-t2)
                        del U
                        del U1
                        
                    chip.img = chip.img + tmp_img*sum_img/np.sum(tmp_img)
                    del tmp_img
                    gc.collect()

        # pixSize = np.rad2deg(self.pixsize*1e-3/28)*3600  #set psf pixsize
        #
        # # assert self.iccd == int(chip.getChipLabel(chipID=chip.chipID)), 'ERROR: self.iccd != chip.chipID'
        # twave = self._findWave(bandpass)
        # if twave == -1:
        #     print("!!!PSF bandpass does not match.")
        #     exit()
        # PSFMat = self.psfMat[twave]
        # cen_col= self.cen_col[twave]
        # cen_row= self.cen_row[twave]
        #
        # px = (pos_img.x - chip.cen_pix_x)*0.01
        # py = (pos_img.y - chip.cen_pix_y)*0.01
        # if findNeighMode == 'treeFind':
        #     imPSF = psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, PSFCentroidWgt=True)
        # if findNeighMode == 'hoclistFind':
        #     assert(self.hoc != 0), 'hoclist should be built correctly!'
        #     imPSF = psfMaker_IDW(px, py, PSFMat, cen_col, cen_row, IDWindex=2, OnlyNeighbors=True, hoc=self.hoc[twave], hoclist=self.hoclist[twave], PSFCentroidWgt=True)
        #
        # ############TEST: START
        # TestGaussian = False
        # if TestGaussian:
        #     gsx  = galsim.Gaussian(sigma=0.04)
        #     #pointing_pa = -23.433333
        #     imPSF= gsx.shear(g1=0.8, g2=0.).rotate(0.*galsim.degrees).drawImage(nx = 256, ny=256, scale=pixSize).array
        # ############TEST: END
        #
        # if galsimGSObject:
        #     imPSFt = np.zeros([257,257])
        #     imPSFt[0:256, 0:256] = imPSF
        #     # imPSFt[120:130, 0:256] = 1.
        #
        #     img = galsim.ImageF(imPSFt, scale=pixSize)
        #     gsp = galsim.GSParams(folding_threshold=folding_threshold)
        #     ############TEST: START
        #     # Use sheared PSF to test the PSF orientation
        #     # self.psf = galsim.InterpolatedImage(img, gsparams=gsp).shear(g1=0.8, g2=0.)
        #     ############TEST: END
        #     self.psf = galsim.InterpolatedImage(img, gsparams=gsp)
        #     wcs = chip.img.wcs.local(pos_img)
        #     scale = galsim.PixelScale(0.074)
        #     self.psf = wcs.toWorld(scale.toImage(self.psf), image_pos=(pos_img))
        #
        #     # return self.PSFspin(x=px/0.01, y=py/0.01)
        #     return self.psf, galsim.Shear(e=0., beta=(np.pi/2)*galsim.radians)
        # return imPSF
    #
    # def PSFspin(self, x, y):
    #     """
    #     The PSF profile at a given image position relative to the axis center
    #
    #     Parameters:
    #     theta : spin angles in a given exposure in unit of [arcsecond]
    #     dx, dy: relative position to the axis center in unit of [pixels]
    #
    #     Return:
    #     Spinned PSF: g1, g2 and axis ratio 'a/b'
    #     """
    #     a2Rad = np.pi/(60.0*60.0*180.0)
    #
    #     ff = self.sigGauss * 0.107 * (1000.0/10.0) # in unit of [pixels]
    #     rc = np.sqrt(x*x + y*y)
    #     cpix = rc*(self.sigSpin*a2Rad)
    #
    #     beta = (np.arctan2(y,x) + np.pi/2)
    #     ell = cpix**2/(2.0*ff**2+cpix**2)
    #     qr = np.sqrt((1.0+ell)/(1.0-ell))
    #     PSFshear = galsim.Shear(e=ell, beta=beta*galsim.radians)
    #     return self.psf.shear(PSFshear), PSFshear


if __name__ == '__main__':
    configfn = '/Users/zhangxin/Work/SlitlessSim/CSST_SIM/CSST_new_sim/csst-simulation/config/config_C6_dev.yaml'
    with open(configfn, "r") as stream:
        try:
            config = yaml.safe_load(stream)
            for key, value in config.items():
                print(key + " : " + str(value))
        except yaml.YAMLError as exc:
            print(exc)
    chip = Chip(chipID=1, config=config)
    filter_id, filter_type = chip.getChipFilter()
    filt = Filter(filter_id=filter_id,
                  filter_type=filter_type,
                  filter_param=FilterParam())

    psf_i = PSFInterpSLS(
        chip, filt, PSF_data_prefix="/Volumes/EAGET/CSST_PSF_data/SLS_PSF_PCA_fp/")
    pos_img = galsim.PositionD(x=25155, y=-22060)
    psf_im = psf_i.get_PSF(chip, pos_img=pos_img, g_order='1')