from numpy.random import Generator, PCG64
from ObservationSim.Instrument.Chip import ChipUtils as chip_utils
from ObservationSim.Instrument.Chip import Effects

def apply_PRNU(self, chip, filt, tel, pointing, catalog, obs_param):
    chip.img *= chip.prnu_img
    if self.overall_config["output_setting"]["prnu_output"] == True:
        chip.prnu_img.write("%s/FlatImg_PRNU_%s.fits" % (self.chip_output.subdir, str(chip.chipID).rjust(2, '0')))
    return chip, filt, tel, pointing

def add_poisson_and_dark(self, chip, filt, tel, pointing, catalog, obs_param):
    # Add dark current & Poisson noise
    # Get exposure time
    if (obs_param) and ("exptime" in obs_param) and (obs_param["exptime"] is not None):
        exptime = obs_param["exptime"]
    else:
        exptime = pointing.exp_time
    
    if obs_param["add_dark"] == True:
        chip.img, _ = chip_utils.add_poisson(img=chip.img,
                                            chip=chip,
                                            exptime=pointing.exp_time,
                                            poisson_noise=chip.poisson_noise,
                                            InputDark=None)
    else:
        chip.img, _ = chip_utils.add_poisson(img=chip.img, 
                                            chip=self, 
                                            exptime=exptime, 
                                            poisson_noise=chip.poisson_noise, 
                                            dark_noise=0.)
    return chip, filt, tel, pointing

def add_detector_defects(self, chip, filt, tel, pointing, catalog, obs_param):
    # Add Hot Pixels or/and Dead Pixels
    rgbadpix = Generator(PCG64(int(self.overall_config["random_seeds"]["seed_defective"]+chip.chipID)))
    badfraction = 5E-5*(rgbadpix.random()*0.5+0.7)
    chip.img = Effects.DefectivePixels(
        chip.img, 
        IfHotPix=obs_param["hot_pixels"], 
        IfDeadPix=obs_param["dead_pixels"],
        fraction=badfraction, 
        seed=self.overall_config["random_seeds"]["seed_defective"]+chip.chipID, biaslevel=0)
    # Apply Bad columns 
    if obs_param["bad_columns"] == True:
        chip.img = Effects.BadColumns(chip.img, 
                                     seed=self.overall_config["random_seeds"]["seed_badcolumns"],  
                                     chipid=chip.chipID)
    return chip, filt, tel, pointing

def add_nonlinearity(self, chip, filt, tel, pointing, catalog, obs_param):
    self.chip_output.Log_info("  Applying Non-Linearity on the chip image")
    chip.img = Effects.NonLinearity(GSImage=chip.img, 
                                   beta1=5.e-7, 
                                   beta2=0)
    return chip, filt, tel, pointing

def add_blooming(self, chip, filt, tel, pointing, catalog, obs_param):
    self.chip_output.Log_info("  Applying CCD Saturation & Blooming")
    chip.img = Effects.SaturBloom(GSImage=chip.img, 
                                 nsect_x=1, 
                                 nsect_y=1, 
                                 fullwell=int(chip.full_well))
    return chip, filt, tel, pointing

def add_bias(self, chip, filt, tel, pointing, catalog, obs_param):
    self.chip_output.Log_info("  Adding Bias level and 16-channel non-uniformity")
    if obs_param["bias_16channel"] == True:
        chip.img = Effects.AddBiasNonUniform16(chip.img, 
                                              bias_level=float(chip.bias_level), 
                                              nsecy = chip.nsecy, 
                                              nsecx=chip.nsecx, 
                                              seed=self.overall_config["random_seeds"]["seed_biasNonUniform"]+chip.chipID)
    elif obs_param["bias_16channel"] == False:
        chip.img += self.bias_level
    return chip, filt, tel, pointing