import os
import numpy as np
import observation_sim.psf.PSFInterp as PSFInterp
from observation_sim.instruments 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)
print('chip.bound::', chip.bound.xmin, chip.bound.xmax,
      chip.bound.ymin, chip.bound.ymax)

for iobj in range(nobj):
    print("\nget psf for iobj-", iobj, '\t', 'bandpass:', end=" ", flush=True)
    # Setup Position on focalplane
    # try get the PSF at some location (1234, 1234) on the chip
    x, y = imgPos[iobj, :]
    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=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)