import os
import numpy as np
import ObservationSim.PSF.PSFInterp as PSFInterp
from ObservationSim.Instrument import Chip, Filter, FilterParam
import yaml
import galsim
import astropy.io.fits as fitsio

# Setup PATH
SIMPATH = "/share/simudata/CSSOSDataProductsSims/data/CSSTSimImage_C8/testRun_FGS"
config_filename= SIMPATH+"/config_C6_fits.yaml"
cat_filename   = SIMPATH+"/MSC_00000000/MSC_10106100000000_chip_40_filt_FGS.cat"

# Read cat file
catFn = open(cat_filename,"r")
line = catFn.readline()
print(cat_filename,'\n',line)
imgPos = []
chipID = -1
for line in catFn:
    line = line.strip()
    columns = line.split()
    
    if chipID == -1:
        chipID = int(columns[1])
    else:
        assert chipID == int(columns[1])
    ximg = float(columns[3])
    yimg = float(columns[4])
    imgPos.append([ximg, yimg])
imgPos = np.array(imgPos)
nobj = imgPos.shape[0]
print('chipID, nobj::', chipID, nobj)

# Read config file
with open(config_filename, "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)

# Setup Chip
chip = Chip(chipID=chipID, config=config)

for iobj in range(nobj):
    print("\nget psf for iobj-", iobj, '\t', 'bandpass:', end=" ", flush=True)
    # Setup Position
    x, y = imgPos[iobj, :] # try get the PSF at some location (1234, 1234) on the chip
    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)
        bandpass = bandpass_list[i] # say you want to access the PSF for the sub-bandpass at the blue end for that chip
        
        # Get corresponding PSF model
        psf_model = PSFInterp(chip=chip, npsf=100, PSF_data_file=config["psf_setting"]["psf_dir"])
        psf  = psf_model.get_PSF(chip=chip, pos_img=pos_img, bandpass=bandpass, galsimGSObject=False)
        
        if True:
            fn = "psf_{:}.{:}.{:}.fits".format(chipID, iobj, i)
            if fn != None:
                if os.path.exists(fn):
                    os.remove(fn)
            hdu = fitsio.PrimaryHDU()
            hdu.data = psf
            hdu.header.set('pixScale', 5)
            hdu.writeto(fn)