import numpy as np
import galsim
from ObservationSim.Straylight import calculateSkyMap_split_g
from ObservationSim.Instrument import FilterParam

def add_sky_background_sci(self, chip, filt, tel, pointing, catalog, obs_param):
    
    # 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
    
    flat_normal = np.ones_like(chip.img.array)
    if obs_param["flat_fielding"] == True:
        flat_normal = flat_normal * chip.flat_img.array / np.mean(chip.flat_img.array)
    if obs_param["shutter_effect"] == True:
        flat_normal = flat_normal * chip.shutter_img
        flat_normal = np.array(flat_normal, dtype='float32')
        self.updateHeaderInfo(header_flag='ext', keys = ['SHTSTAT'], values = [True])
    else:
        self.updateHeaderInfo(header_flag='ext', keys = ['SHTSTAT','SHTOPEN0','SHTOPEN1','SHTCLOS0','SHTCLOS1'], values = [False,'','','',''])
    
    if obs_param["enable_straylight_model"]:
        # Filter.sky_background, Filter.zodical_spec will be updated
        filt.setFilterStrayLightPixel(
            jtime = pointing.jdt,
            sat_pos = np.array([pointing.sat_x, pointing.sat_y, pointing.sat_z]),
            pointing_radec = np.array([pointing.ra,pointing.dec]),
            sun_pos = np.array([pointing.sun_x, pointing.sun_y, pointing.sun_z]))
        self.chip_output.Log_info("================================================")
        self.chip_output.Log_info("sky background + stray light pixel flux value: %.5f"%(filt.sky_background))
    
    if chip.survey_type == "photometric":
        sky_map = filt.getSkyNoise(exptime = exptime)
        sky_map = sky_map * np.ones_like(chip.img.array) * flat_normal
        sky_map = galsim.Image(array=sky_map)
    else:
        # chip.loadSLSFLATCUBE(flat_fn='flat_cube.fits')
        sky_map = calculateSkyMap_split_g(
                skyMap=flat_normal,
                blueLimit=filt.blue_limit,
                redLimit=filt.red_limit,
                conf=chip.sls_conf,
                pixelSize=chip.pix_scale,
                isAlongY=0,
                flat_cube=chip.flat_cube, 
                zoldial_spec = filt.zodical_spec)
        sky_map = (sky_map + filt.sky_background)*exptime
    
    # sky_map = sky_map * tel.pupil_area * obs_param["exptime"]
    chip.img += sky_map
    return chip, filt, tel, pointing

def add_sky_flat_calibration(self, chip, filt, tel, pointing, catalog, obs_param):

    if not hasattr(self, 'h_ext'):
        _, _ = self.prepare_headers(chip=chip, pointing=pointing)
    chip_wcs = galsim.FitsWCS(header = self.h_ext)

    # 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
    
    skyback_level = obs_param["flat_level"]

    filter_param = FilterParam()
    sky_level_filt = obs_param["flat_level_filt"]
    norm_scaler = skyback_level/exptime / filter_param.param[sky_level_filt][5]

    flat_normal = np.ones_like(chip.img.array)
    if obs_param["flat_fielding"] == True:
        flat_normal = flat_normal * chip.flat_img.array / np.mean(chip.flat_img.array)
    if obs_param["shutter_effect"] == True:
        flat_normal = flat_normal * chip.shutter_img
        flat_normal = np.array(flat_normal, dtype='float32')
        if self.overall_config["output_setting"]["shutter_output"] == True:    # output 16-bit shutter effect image with pixel value <=65535
            shutt_gsimg = galsim.ImageUS(chip.shutter_img*6E4)
            shutt_gsimg.write("%s/ShutterEffect_%s_1.fits" % (self.chip_output.subdir, chip.chipID))
        self.updateHeaderInfo(header_flag='ext', keys = ['SHTSTAT'], values = [True])
    else:
        self.updateHeaderInfo(header_flag='ext', keys = ['SHTSTAT','SHTOPEN0','SHTOPEN1','SHTCLOS0','SHTCLOS1'], values = [False,'','','',''])
    

    if chip.survey_type == "photometric":
        sky_map = flat_normal * np.ones_like(chip.img.array) * norm_scaler * filter_param.param[chip.filter_type][5] / tel.pupil_area * exptime
    elif chip.survey_type == "spectroscopic":
        # flat_normal = np.ones_like(chip.img.array)
        if obs_param["flat_fielding"] == True:
            
            flat_normal = flat_normal * chip.flat_img.array / np.mean(chip.flat_img.array)
        if obs_param["shutter_effect"] == True:
            
            flat_normal = flat_normal * chip.shutter_img
            flat_normal = np.array(flat_normal, dtype='float32')
        sky_map = calculateSkyMap_split_g(
            skyMap=flat_normal,
            blueLimit=filt.blue_limit,
            redLimit=filt.red_limit,
            conf=chip.sls_conf,
            pixelSize=chip.pix_scale,
            isAlongY=0,
            flat_cube=chip.flat_cube)
        sky_map = sky_map * norm_scaler * exptime
    
    chip.img += sky_map
    return chip, filt, tel, pointing

def add_sky_background(self, chip, filt, tel, pointing, catalog, obs_param):
    if not hasattr(self, 'h_ext'):
        _, _ = self.prepare_headers(chip=chip, pointing=pointing)
    chip_wcs = galsim.FitsWCS(header = self.h_ext)

    if "flat_level" not in obs_param or "flat_level_filt" not in obs_param:
        chip, filt, tel, pointing = self.add_sky_background_sci(chip, filt, tel, pointing, catalog, obs_param)
    else:
        if obs_param.get('flat_level') is None or obs_param.get('flat_level_filt')is None:
            chip, filt, tel, pointing = self.add_sky_background_sci(chip, filt, tel, pointing, catalog, obs_param)
        else:
            chip, filt, tel, pointing = self.add_sky_flat_calibration(chip, filt, tel, pointing, catalog, obs_param)
    return chip, filt, tel, pointing