import numpy as np
import observation_sim.psf.PSFInterpSLS as PSFInterpSLS
import observation_sim.psf.PSFInterp as PSFInterp
from observation_sim.instruments import Chip, Filter, FilterParam
import astropy.io.fits as fitsio

from observation_sim.instruments import Chip, FilterParam, Filter
import galsim

# 计算 0级或1级光谱在某个波长位置的PSF, 返回值是过采样的PSF，像元大小有CSST图像像元大小的1/2
# chipId 必须为[1,2,3,4,5,10,21,26,27,28,29,30]
# order 只有 0 或 1
# pos_img 为直接成像在图像上的位置,[x, y]
# wave: 波长，单位A


def get_SLS_PSF(chipID=1, pos_img=[6000, 4000], order=1, wave=8000, sls_psf_dir='/nfsdata/share/CSSOSDataProductsSims/data/SLS_PSF_PCA_fp_cd/'):
    orders = {0: 'B', 1: 'A'}
    chip = Chip(chipID)
    filter_id, filter_type = chip.getChipFilter()
    filt = Filter(
        filter_id=filter_id,
        filter_type=filter_type,
        filter_param=FilterParam())

    psf_model = PSFInterpSLS(chip, filt, PSF_data_prefix=sls_psf_dir)
    bandNo = 1
    for i, brange in enumerate(psf_model.bandranges):
        if wave >= brange[0] and wave < brange[1]:
            bandNo = i + 1
            break

    psf1, _ = psf_model.get_PSF(
        chip, pos_img_local=pos_img, bandNo=bandNo, galsimGSObject=False, g_order=orders[order])
    fn = "psf.chip{:}.order{:}.wave{:}.fits".format(chipID, order, wave)
    hdu = fitsio.PrimaryHDU()
    hdu.data = psf1
    hdu.header.set('pixScale', 5)
    hdu.writeto(fn, overwrite=True)

# 获得成像的PSF,PSF在波长区间分四段采样，返回4段的PSF，每个PSF代表不同band波长范围不同的四段psf，波长从小到大
# chipId 必须为[6,7,8,9,11,12,13,14,15,16,17,18,19,20,22,23,24,25]
# pos_img 为直接成像在图像上的位置,[x, y]


def getPSFImage(chipID=6, pos_img=[6000, 4500], psf_pho_dir="/nfsdata/share/CSSOSDataProductsSims/data/psfcube/set1_dynamic/"):
    chip = Chip(chipID=chipID)
    print('chip.bound::', chip.bound.xmin, chip.bound.xmax,
          chip.bound.ymin, chip.bound.ymax)

    # Setup Position on focalplane
    # try get the PSF at some location (1234, 1234) on the chip
    pos_img = pos_img
    x, y = pos_img
    x = x+chip.bound.xmin
    y = y+chip.bound.ymin

    pos_img = galsim.PositionD(x, y)

    # Setup sub-bandpass
    # (There are 4 sub-bandpasses for each PSF sample)
    filter_param = FilterParam()
    filter_id, filter_type = chip.getChipFilter()
    filt = Filter(
        filter_id=filter_id,
        filter_type=filter_type,
        filter_param=filter_param,
        ccd_bandpass=chip.effCurve)
    bandpass_list = filt.bandpass_sub_list
    for i in range(len(bandpass_list)):
        print(i, end=" ", flush=True)
        # say you want to access the PSF for the sub-bandpass at the blue end for that chip
        bandpass = bandpass_list[i]

        # Get corresponding PSF model
        psf_model = PSFInterp(chip=chip, npsf=100,
                              PSF_data_file=psf_pho_dir)
        psf = psf_model.get_PSF(
            chip=chip, pos_img=pos_img, bandpass=bandpass, galsimGSObject=False)
        fn = "psf_chip{:}.wave{:}.fits".format(chipID, i)
        hdu = fitsio.PrimaryHDU()
        hdu.data = psf
        hdu.header.set('pixScale', 5)
        hdu.writeto(fn, overwrite=True)


if __name__ == "__main__":
    chipid = 2
    pos_img = [6000, 4500]
    order = 1
    w = 5000
    get_SLS_PSF(chipID=chipid, pos_img=pos_img, order=order, wave=w)

    chipid = 7
    pos_img = [6000, 4500]
    getPSFImage(chipID=chipid, pos_img=pos_img)