new sim: modified ways to add detector effects: e.g. brighter-fatter etc.

from ObservationSim.Instrument.FocalPlane import FocalPlane

from ObservationSim.Config.Header import generatePrimaryHeader, generateExtensionHeader

from ObservationSim.Instrument._util import rotate_conterclockwise

from ObservationSim.Instrument.Chip import ChipUtils as chip_utils

try:

    import importlib.resources as pkg_resources

class Chip(FocalPlane):

    def __init__(self, chipID, ccdEffCurve_dir=None, CRdata_dir=None, sls_dir=None, config=None, treering_func=None, logger=None):

        # Get focal plane (instance of paraent class) info

        # TODO: use chipID to config individual chip?

        super().__init__()

        # self.npix_x = 9216

        # self.npix_y = 9232

        # self.pix_scale  = 0.074 # pixel scale

        self.nsecy = 2

        self.nsecx = 8

        self.gain_channel = np.ones(self.nsecy* self.nsecx)

        self.filter_id, self.filter_type = self.getChipFilter()

        self.survey_type = self._getSurveyType()

        # [TODO]

        if self.filter_type != "FGS":

            self._getChipRowCol()

        # Set the relavent specs for FGS detectors

        # [TODO]

        # Set the relavent specs for detectors

        try:

            with pkg_resources.files('ObservationSim.Instrument.data.ccd').joinpath("chip_definition.json") as chip_definition:

                with open(chip_definition, "r") as f:

                    chip_dict = json.load(f)[str(self.chipID)]

        for key in chip_dict:

            setattr(self, key, chip_dict[key])

        if self.filter_type == "FGS":

            if ("field_dist" in config) and (config["ins_effects"]["field_dist"]) == False:

                self.fdModel = None

                self.fdModel = None

            else:

                try:

                    # with pkg_resources.files('ObservationSim.Instrument.data.field_distortion').joinpath("FieldDistModelGlobal_mainFP_v1.0.pickle") as field_distortion:

                    with pkg_resources.files('ObservationSim.Instrument.data.field_distortion').joinpath("FieldDistModel_v2.0.pickle") as field_distortion:

                        with open(field_distortion, "rb") as f:

                            self.fdModel = pickle.load(f)

        # Get boundary (in pix)

        self.bound = self.getChipLim()

        self.ccdEffCurve_dir = ccdEffCurve_dir

        self.CRdata_dir = CRdata_dir

        slsconfs = self.getChipSLSConf()

        slsconfs = chip_utils.getChipSLSConf(chipID=self.chipID)

        if np.size(slsconfs) == 1:

            try:

                with pkg_resources.files('ObservationSim.Instrument.data.sls_conf').joinpath(slsconfs) as conf_path:

        if filter_type in ['NUV', 'u', 'GU']: filename = 'UV0.txt'

        if filter_type in ['g', 'r', 'GV', 'FGS']: filename = 'Astro_MB.txt' # TODO, need to switch to the right efficiency curvey for FGS CMOS

        if filter_type in ['i', 'z', 'y', 'GI']: filename = 'Basic_NIR.txt'

        # Mirror efficiency:

        # if filter_type == 'NUV': mirror_eff = 0.54

        # if filter_type == 'u': mirror_eff = 0.68

        # if filter_type in ['g', 'r', 'i', 'z', 'y']: mirror_eff = 0.8

        # if filter_type in ['GU', 'GV', 'GI']: mirror_eff = 1. # Not sure if this is right

        # path = os.path.join(self.ccdEffCurve_dir, filename)

        # table = Table.read(path, format='ascii')

        try:

            with pkg_resources.files('ObservationSim.Instrument.data.ccd').joinpath(filename) as ccd_path:

                table = Table.read(ccd_path, format='ascii')

            with pkg_resources.path('ObservationSim.Instrument.data', "wfc-cr-attachpixel.dat") as cr_path:

                self.attachedSizes = np.loadtxt(cr_path)

    def getChipFilter(self, chipID=None, filter_layout=None):

    def loadSLSFLATCUBE(self, flat_fn='flat_cube.fits'):

        try:

            with pkg_resources.files('ObservationSim.Instrument.data').joinpath(flat_fn) as data_path:

                flat_fits = fits.open(data_path, ignore_missing_simple=True)

        except AttributeError:

            with pkg_resources.path('ObservationSim.Instrument.data', flat_fn) as data_path:

                flat_fits = fits.open(data_path, ignore_missing_simple=True)

        fl = len(flat_fits)

        fl_sh = flat_fits[0].data.shape

        assert fl == 4, 'FLAT Field Cube is Not 4 layess!!!!!!!'

        self.flat_cube = np.zeros([fl, fl_sh[0], fl_sh[1]])

        for i in np.arange(0, fl, 1):

            self.flat_cube[i] = flat_fits[i].data

    def getChipFilter(self, chipID=None):

        """Return the filter index and type for a given chip #(chipID)

        """

        filter_type_list = ["NUV","u", "g", "r", "i","z","y","GU", "GV", "GI", "FGS"]

        if filter_layout is not None:

            return filter_layout[chipID][0], filter_layout[chipID][1]

        if chipID == None:

            chipID = self.chipID

        Returns:

            A galsim BoundsD object

        """

        if ((chipID is not None) and (int(chipID) <= 30)) or (self.chipID <= 30):

            # [TODO]

            if chipID == None:

                chipID = self.chipID

                rowID, colID = self.rowID, self.colID

            else:

                rowID, colID = self.getChipRowCol(chipID)

            gx1, gx2 = self.npix_gap_x

            gy = self.npix_gap_y

            # xlim of a given CCD chip

            xrem = 2*(colID - 1) - (self.nchip_x - 1)

            xcen = (self.npix_x//2 + gx1//2) * xrem

            if chipID >= 26 or chipID == 21:

                xcen = (self.npix_x//2 + gx1//2) * xrem - (gx2-gx1)

            if chipID <= 5 or chipID == 10:

                xcen = (self.npix_x//2 + gx1//2) * xrem + (gx2-gx1)

            nx0 = xcen - self.npix_x//2 + 1

            nx1 = xcen + self.npix_x//2

            # ylim of a given CCD chip

            yrem = (rowID - 1) - self.nchip_y // 2

            ycen = (self.npix_y + gy) * yrem

            ny0 = ycen - self.npix_y//2 + 1

            ny1 = ycen + self.npix_y//2

            return galsim.BoundsD(nx0-1, nx1-1, ny0-1, ny1-1)

        else:

        xmin, xmax, ymin, ymax = 1e10, -1e10, 1e10, -1e10

        xcen = self.x_cen / self.pix_size

        ycen = self.y_cen / self.pix_size

        noise = self.dark_noise * exptime + self.read_noise**2

        return noise

    def getChipSLSConf(self):

        confFile = ''

        if self.chipID == 1: confFile = ['CSST_GI2.conf', 'CSST_GI1.conf']

        if self.chipID == 2: confFile = ['CSST_GV4.conf', 'CSST_GV3.conf']

        if self.chipID == 3: confFile = ['CSST_GU2.conf', 'CSST_GU1.conf']

        if self.chipID == 4: confFile = ['CSST_GU4.conf', 'CSST_GU3.conf']

        if self.chipID == 5: confFile = ['CSST_GV2.conf', 'CSST_GV1.conf']

        if self.chipID == 10: confFile = ['CSST_GI4.conf', 'CSST_GI3.conf']

        if self.chipID == 21: confFile = ['CSST_GI6.conf', 'CSST_GI5.conf']

        if self.chipID == 26: confFile = ['CSST_GV8.conf', 'CSST_GV7.conf']

        if self.chipID == 27: confFile = ['CSST_GU6.conf', 'CSST_GU5.conf']

        if self.chipID == 28: confFile = ['CSST_GU8.conf', 'CSST_GU7.conf']

        if self.chipID == 29: confFile = ['CSST_GV6.conf', 'CSST_GV5.conf']

        if self.chipID == 30: confFile = ['CSST_GI8.conf', 'CSST_GI7.conf']

        return confFile

    def generateHeader(self, ra_cen, dec_cen, img_rot, im_type, pointing_ID, exptime=150., timestamp = 1621915200):

        datetime_obs = datetime.utcfromtimestamp(timestamp)

        date_obs = datetime_obs.strftime("%y%m%d")

        time_obs = datetime_obs.strftime("%H%M%S")

        h_prim = generatePrimaryHeader(

            xlen=self.npix_x, 

            ylen=self.npix_y, 

            pointNum = str(pointing_ID),

            ra=ra_cen, 

            dec=dec_cen, 

            pixel_scale=self.pix_scale,

            date=date_obs,

            time_obs=time_obs,

            im_type = im_type,

            exptime=exptime,

            chip_name=str(self.chipID).rjust(2, '0')

            )

        h_ext = generateExtensionHeader(

            chip=self,

            xlen=self.npix_x, 

            ylen=self.npix_y, 

            ra=ra_cen, 

            dec=dec_cen,

            pa=img_rot.deg, 

            gain=self.gain, 

            readout=self.read_noise, 

            dark=self.dark_noise, 

            saturation=90000, 

            pixel_scale=self.pix_scale, 

            pixel_size=self.pix_size,

            xcen=self.x_cen,

            ycen=self.y_cen,

            extName='SCI',

            timestamp = timestamp,

            exptime = exptime,

            readoutTime = 40.)

        return h_prim, h_ext

    def outputCal(self, img, ra_cen, dec_cen, img_rot, im_type, pointing_ID, output_dir, exptime=150., timestamp = 1621915200):

        h_prim, h_ext = self.generateHeader(

            ra_cen=ra_cen,

            dec_cen=dec_cen,

            img_rot=img_rot,

            im_type=im_type,

            pointing_ID=pointing_ID,

            exptime=exptime,

            timestamp = timestamp)

        hdu1 = fits.PrimaryHDU(header=h_prim)

        hdu1.add_checksum()

        hdu1.header.comments['CHECKSUM'] = 'HDU checksum'

        hdu1.header.comments['DATASUM'] = 'data unit checksum'

        hdu2 = fits.ImageHDU(img.array, header=h_ext)

        hdu2.add_checksum()

        hdu2.header.comments['XTENSION'] = 'extension type'

        hdu2.header.comments['CHECKSUM'] = 'HDU checksum'

        hdu2.header.comments['DATASUM'] = 'data unit checksum'

        hdu1 = fits.HDUList([hdu1, hdu2])

        fname = os.path.join(output_dir, h_prim['FILENAME']+'.fits')

        hdu1.writeto(fname, output_verify='ignore', overwrite=True)

    def addEffects(self, config, img, chip_output, filt, ra_cen, dec_cen, img_rot, exptime=150., pointing_ID=0, timestamp_obs=1621915200, pointing_type='MS', sky_map=None, tel=None, logger=None):

        # Set random seeds

        SeedGainNonuni=int(config["random_seeds"]["seed_gainNonUniform"])

        SeedBiasNonuni=int(config["random_seeds"]["seed_biasNonUniform"])

        SeedRnNonuni = int(config["random_seeds"]["seed_rnNonUniform"])

        self.logger = logger

        # Get Poisson noise generator

        seed = int(config["random_seeds"]["seed_poisson"]) + pointing_ID*30 + self.chipID

        rng_poisson = galsim.BaseDeviate(seed)

        poisson_noise = galsim.PoissonNoise(rng_poisson, sky_level=0.)

        rng_poisson, poisson_noise = chip_utils.get_poisson(

            seed=int(config["random_seeds"]["seed_poisson"]) + pointing_ID*30 + self.chipID, sky_level=0.)

        # Add sky background

        if sky_map is None:

            sky_map = filt.getSkyNoise(exptime=exptime)

            sky_map = sky_map * np.ones_like(img.array)

            sky_map = galsim.Image(array=sky_map)

            # Apply Poisson noise to the sky map

            # (NOTE): only for photometric chips

            # since it utilize the photon shooting

            # to draw stamps

            if self.survey_type == "photometric":

                sky_map.addNoise(poisson_noise)

        elif img.array.shape != sky_map.shape:

            raise ValueError("The shape img and sky_map must be equal.")

        elif tel is not None: # If sky_map is given in flux

            sky_map = sky_map * tel.pupil_area * exptime

        if config["ins_effects"]["add_back"] == True:

            img += sky_map

            img, sky_map = chip_utils.add_sky_background(img=img, filt=filt, exptime=exptime, sky_map=sky_map, tel=tel)

            del sky_map

        # Apply flat-field large scale structure for one chip

        if config["ins_effects"]["flat_fielding"] == True:

            if self.logger is not None:

                self.logger.info("  Creating and applying Flat-Fielding")

                msg = str(img.bounds)

                self.logger.info(msg)

            else:

                print("  Creating and applying Flat-Fielding", flush=True)

                print(img.bounds, flush=True)

            flat_img = effects.MakeFlatSmooth(

                img.bounds, 

                int(config["random_seeds"]["seed_flat"]))

            flat_normal = flat_img / np.mean(flat_img.array)

            chip_utils.log_info(msg="  Creating and applying Flat-Fielding", logger=self.logger)

            chip_utils.log_info(msg=str(img.bounds), logger=self.logger)

            flat_img, flat_normal = chip_utils.get_flat(img=img, seed=int(config["random_seeds"]["seed_flat"]))

            if self.survey_type == "photometric":

                img *= flat_normal

            del flat_normal

        # Apply Shutter-effect for one chip

        if config["ins_effects"]["shutter_effect"] == True:

            if self.logger is not None:

                self.logger.info("  Apply shutter effect")

            else:

                print("  Apply shutter effect", flush=True)

            chip_utils.log_info(msg="  Apply shutter effect", logger=self.logger)

            shuttimg = effects.ShutterEffectArr(img, t_shutter=1.3, dist_bearing=735, dt=1E-3)    # shutter effect normalized image for this chip

            if self.survey_type == "photometric":

                img *= shuttimg

                del shutt_gsimg

            del shuttimg

        # Add Poisson noise to the resulting images

        # (NOTE): this can only applied to the slitless image

        # since it dose not use photon shooting to draw stamps

        if self.survey_type == "spectroscopic":

            img.addNoise(poisson_noise)

        # # Add Poisson noise to the resulting images

        # # (NOTE): this can only applied to the slitless image

        # # since it dose not use photon shooting to draw stamps

        # if self.survey_type == "spectroscopic":

        #     img.addNoise(poisson_noise)

        # Add cosmic-rays

        if config["ins_effects"]["cosmic_ray"] == True and pointing_type=='MS':

            if self.logger is not None:

                self.logger.info(("  Adding Cosmic-Ray"))

            else:

                print("  Adding Cosmic-Ray", flush=True)

            cr_map, cr_event_num = effects.produceCR_Map(

                xLen=self.npix_x, yLen=self.npix_y, 

                exTime=exptime+0.5*self.readout_time, 

                cr_pixelRatio=0.003*(exptime+0.5*self.readout_time)/600.,

                gain=self.gain, 

                attachedSizes=self.attachedSizes,

                seed=SeedCosmicRay+pointing_ID*30+self.chipID)   # seed: obj-imaging:+0; bias:+1; dark:+2; flat:+3;

            img += cr_map

            cr_map[cr_map > 65535] = 65535

            cr_map[cr_map < 0] = 0

            crmap_gsimg = galsim.Image(cr_map, dtype=np.uint16)

            del cr_map

            # crmap_gsimg.write("%s/CosmicRay_%s_1.fits" % (chip_output.subdir, self.chipID))

            # crmap_gsimg.write("%s/CosmicRay_%s.fits" % (chip_output.subdir, self.chipID))

            # datetime_obs = datetime.utcfromtimestamp(timestamp_obs)

            # date_obs = datetime_obs.strftime("%y%m%d")

            # time_obs = datetime_obs.strftime("%H%M%S")

            self.outputCal(

            chip_utils.log_info(msg="  Adding Cosmic-Ray", logger=self.logger)

            img, crmap_gsimg, cr_event_num = chip_utils.add_cosmic_rays(img=img, chip=self, exptime=exptime, 

                                                    seed=SeedCosmicRay+pointing_ID*30+self.chipID)

            chip_utils.outputCal(

                chip=self,

                img=crmap_gsimg,

                ra_cen=ra_cen,

                dec_cen=dec_cen,

        # Apply PRNU effect and output PRNU flat file:

        if config["ins_effects"]["prnu_effect"] == True:

            if self.logger is not None:

                self.logger.info("  Applying PRNU effect")

            else:

                print("  Applying PRNU effect", flush=True)

            prnu_img = effects.PRNU_Img(

                xsize=self.npix_x, 

                ysize=self.npix_y, 

                sigma=0.01, 

            chip_utils.log_info(msg="  Applying PRNU effect", logger=self.logger)

            img, prnu_img = chip_utils.add_PRNU(img=img, chip=self, 

                                seed=int(config["random_seeds"]["seed_prnu"]+self.chipID))

            img *= prnu_img

            if config["output_setting"]["prnu_output"] == True:

                prnu_img.write("%s/FlatImg_PRNU_%s.fits" % (chip_output.subdir,self.chipID))

            if config["output_setting"]["flat_output"] == False:

                del prnu_img

        # Add dark current

        # # Add dark current

        # if config["ins_effects"]["add_dark"] == True:

        #     dark_noise = galsim.DeviateNoise(galsim.PoissonDeviate(rng_poisson, self.dark_noise*(exptime+0.5*self.readout_time)))

        #     img.addNoise(dark_noise)

        # Add dark current & Poisson noise

        if config["ins_effects"]["add_dark"] == True:

            dark_noise = galsim.DeviateNoise(galsim.PoissonDeviate(rng_poisson, self.dark_noise*(exptime+0.5*self.readout_time)))

            img.addNoise(dark_noise)

            img, _ = chip_utils.add_poisson(img=img, chip=self, exptime=exptime, poisson_noise=poisson_noise)

        else:

            img, _ = chip_utils.add_poisson(img=img, chip=self, exptime=exptime, poisson_noise=poisson_noise, dark_noise=0.)

        # Add diffusion & brighter-fatter effects

        if config["ins_effects"]["bright_fatter"] == True:

            img = chip_utils.add_brighter_fatter(img=img)

        # Add Hot Pixels or/and Dead Pixels

        rgbadpix = Generator(PCG64(int(SeedDefective+self.chipID)))

        # Apply Nonlinearity on the chip image

        if config["ins_effects"]["non_linear"] == True:

            if self.logger is not None:

                self.logger.info("  Applying Non-Linearity on the chip image")

            else:

                print("  Applying Non-Linearity on the chip image", flush=True)

            chip_utils.log_info(msg="  Applying Non-Linearity on the chip image", logger=self.logger)

            img = effects.NonLinearity(GSImage=img, beta1=5.e-7, beta2=0)

        # Apply CCD Saturation & Blooming

        if config["ins_effects"]["saturbloom"] == True:

            if self.logger is not None:

                self.logger.info("  Applying CCD Saturation & Blooming")

            else:

                print("  Applying CCD Saturation & Blooming")

            chip_utils.log_info(msg="  Applying CCD Saturation & Blooming", logger=self.logger)

            img = effects.SaturBloom(GSImage=img, nsect_x=1, nsect_y=1, fullwell=fullwell)

        # Apply CTE Effect

        if config["ins_effects"]["cte_trail"] == True:

            if self.logger is not None:

                self.logger.info("  Apply CTE Effect")

            else:

                print("  Apply CTE Effect")

            chip_utils.log_info(msg="  Apply CTE Effect", logger=self.logger)

            img = effects.CTE_Effect(GSImage=img, threshold=27)

        # Add Bias level

        if config["ins_effects"]["add_bias"] == True:

            if self.logger is not None:

                self.logger.info("  Adding Bias level and 16-channel non-uniformity")

            else:

                print("  Adding Bias level and 16-channel non-uniformity")

            chip_utils.log_info(msg="  Adding Bias level and 16-channel non-uniformity", logger=self.logger)

            if config["ins_effects"]["bias_16channel"] == True:

                img = effects.AddBiasNonUniform16(img, 

                    bias_level=float(self.bias_level), 

            img.addNoise(readout_noise)

        # Apply Gain & Quantization

        if self.logger is not None:

            self.logger.info("  Applying Gain (and 16 channel non-uniformity) & Quantization")

        else:

            print("  Applying Gain (and 16 channel non-uniformity) & Quantization", flush=True)

        chip_utils.log_info(msg="  Applying Gain (and 16 channel non-uniformity) & Quantization", logger=self.logger)

        if config["ins_effects"]["gain_16channel"] == True:

            img, self.gain_channel = effects.ApplyGainNonUniform16(

                img, gain=self.gain, 

                BiasCombImg.replaceNegative(replace_value=0)

                BiasCombImg.quantize()

                BiasCombImg = galsim.ImageUS(BiasCombImg)

                # BiasCombImg.write("%s/BiasImg_%s_%s_%s.fits" % (chip_output.subdir, BiasTag, self.chipID, i+1))

                # datetime_obs = datetime.utcfromtimestamp(timestamp_obs)

                # date_obs = datetime_obs.strftime("%y%m%d")

                # time_obs = datetime_obs.strftime("%H%M%S")

                timestamp_obs += 10 * 60

                self.outputCal(

                chip_utils.outputCal(

                    chip=self,

                    img=BiasCombImg,

                    ra_cen=ra_cen,

                    dec_cen=dec_cen,

                FlatCombImg.replaceNegative(replace_value=0)

                FlatCombImg.quantize()

                FlatCombImg = galsim.ImageUS(FlatCombImg)

                # FlatCombImg.write("%s/FlatImg_%s_%s_%s.fits" % (chip_output.subdir, FlatTag, self.chipID, i+1))

                # datetime_obs = datetime.utcfromtimestamp(timestamp_obs)

                # date_obs = datetime_obs.strftime("%y%m%d")

                # time_obs = datetime_obs.strftime("%H%M%S")

                timestamp_obs += 10 * 60

                self.outputCal(

                chip_utils.outputCal(

                    chip=self,

                    img=FlatCombImg,

                    ra_cen=ra_cen,

                    dec_cen=dec_cen,

                    cr_map[cr_map < 0] = 0

                    crmap_gsimg = galsim.Image(cr_map, dtype=np.uint16)

                    del cr_map

                    # datetime_obs = datetime.utcfromtimestamp(timestamp_obs)

                    # date_obs = datetime_obs.strftime("%y%m%d")

                    # time_obs = datetime_obs.strftime("%H%M%S")

                    self.outputCal(

                    chip_utils.outputCal(

                        chip=self,

                        img=crmap_gsimg,

                        ra_cen=ra_cen,

                        dec_cen=dec_cen,

                DarkCombImg.replaceNegative(replace_value=0)

                DarkCombImg.quantize()

                DarkCombImg = galsim.ImageUS(DarkCombImg)

                # DarkCombImg.write("%s/DarkImg_%s_%s_%s.fits" % (chip_output.subdir, DarkTag, self.chipID, i+1))

                # datetime_obs = datetime.utcfromtimestamp(timestamp_obs)

                # date_obs = datetime_obs.strftime("%y%m%d")

                # time_obs = datetime_obs.strftime("%H%M%S")

                timestamp_obs += 10 * 60

                self.outputCal(

                chip_utils.outputCal(

                    chip=chip,

                    img=DarkCombImg,

                    ra_cen=ra_cen,

                    dec_cen=dec_cen,

        #     del sub_img

        return img

    def loadSLSFLATCUBE(self, flat_fn='flat_cube.fits'):

        from astropy.io import fits

        try:

            with pkg_resources.files('ObservationSim.Instrument.data').joinpath(flat_fn) as data_path:

                flat_fits = fits.open(data_path, ignore_missing_simple=True)

        except AttributeError:

            with pkg_resources.path('ObservationSim.Instrument.data', flat_fn) as data_path:

                flat_fits = fits.open(data_path, ignore_missing_simple=True)

        fl = len(flat_fits)

        fl_sh = flat_fits[0].data.shape

        assert fl == 4, 'FLAT Field Cube is Not 4 layess!!!!!!!'

        self.flat_cube = np.zeros([fl, fl_sh[0], fl_sh[1]])

        for i in np.arange(0, fl, 1):

            self.flat_cube[i] = flat_fits[i].data