import os
import galsim
import numpy as np
from astropy.io import fits
from ObservationSim.Instrument.Chip import ChipUtils as chip_utils
from ObservationSim.Instrument.Chip import Effects

def add_prescan_overscan(self, chip, filt, tel, pointing, catalog, obs_param):
    self.chip_output.Log_info("Apply pre/over-scan")
    chip.img = chip_utils.AddPreScan(GSImage=chip.img,
                                    pre1=chip.prescan_x,
                                    pre2=chip.prescan_y,
                                    over1=chip.overscan_x,
                                    over2=chip.overscan_y)
    return chip, filt, tel, pointing

def add_readout_noise(self, chip, filt, tel, pointing, catalog, obs_param):
    seed = int(self.overall_config["random_seeds"]["seed_readout"]) + pointing.id*30 + chip.chipID
    rng_readout = galsim.BaseDeviate(seed)
    readout_noise = galsim.GaussianNoise(rng=rng_readout, sigma=chip.read_noise)
    chip.img.addNoise(readout_noise)
    return chip, filt, tel, pointing

def apply_gain(self, chip, filt, tel, pointing, catalog, obs_param):
    self.chip_output.Log_info("  Applying Gain")
    if obs_param["gain_16channel"] == True:
        chip.img, chip.gain_channel = Effects.ApplyGainNonUniform16(chip.img, 
                                                                   gain=chip.gain, 
                                                                   nsecy = chip.nsecy, 
                                                                   nsecx=chip.nsecx, 
                                                                   seed=self.overall_config["random_seeds"]["seed_gainNonUniform"]+chip.chipID)
    elif obs_param["gain_16channel"] == False:
        chip.img /= chip.gain
    return chip, filt, tel, pointing

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

    if not hasattr(self, 'h_ext'):
        _, _ = self.prepare_headers(chip=chip, pointing=pointing)
        self.updateHeaderInfo(header_flag='ext', keys = ['SHTSTAT','SHTOPEN0','SHTOPEN1','SHTCLOS0','SHTCLOS1'], values = [False,'','','',''])

    if obs_param["format_output"] == True:
        self.chip_output.Log_info("  Apply 1*16 format")
        chip.img = chip_utils.formatOutput(GSImage=chip.img)
        chip.nsecy = 1
        chip.nsecx = 16

    chip.img.array[chip.img.array > 65535] = 65535
    chip.img.replaceNegative(replace_value=0)
    chip.img.quantize()
    chip.img = galsim.Image(chip.img.array, dtype=np.uint16)
    fname = os.path.join(self.chip_output.subdir, self.h_prim['FILENAME'] + '.fits')

    f_name_size = 68
    if(len(self.h_prim['FILENAME'])>f_name_size):
        self.updateHeaderInfo(header_flag='prim', keys = ['FILENAME'], values = [self.h_prim['FILENAME'][0:f_name_size]])
        
    
    
    hdu1 = fits.PrimaryHDU(header=self.h_prim)
    hdu1.add_checksum()
    hdu1.header.comments['CHECKSUM'] = 'HDU checksum'
    hdu1.header.comments['DATASUM'] = 'data unit checksum'
    self.updateHeaderInfo(header_flag='ext', keys = ['DATASECT'], values = [str(chip.img.array.shape[1])+'x'+str(chip.img.array.shape[0])])
    hdu2 = fits.ImageHDU(chip.img.array, header=self.h_ext)
    hdu2.add_checksum()
    hdu2.header.comments['XTENSION'] = 'extension type'
    hdu2.header.comments['CHECKSUM'] = 'HDU checksum'
    hdu2.header.comments['DATASUM'] = 'data unit checksum'
    hdu2.header.comments["XTENSION"] = "image extension"
    hdu1 = fits.HDUList([hdu1, hdu2])
    hdu1.writeto(fname, output_verify='ignore', overwrite=True)

    return chip, filt, tel, pointing