import numpy as np
from ObservationSim.MockObject import FlatLED
import galsim

from astropy.time import Time
from datetime import datetime, timezone

import gc

def add_LED_Flat(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)
    pf_map = np.zeros_like(chip.img.array)
    if obs_param["LED_TYPE"] is not None:
        if len(obs_param["LED_TYPE"]) != 0:
            print("LED OPEN--------")

            led_obj = FlatLED(chip, filt)
            led_flat, ledstat, letts = led_obj.drawObj_LEDFlat(led_type_list=obs_param["LED_TYPE"], exp_t_list=obs_param["LED_TIME"])
            pf_map = led_flat
            self.updateHeaderInfo(header_flag='ext', keys = ['LEDSTAT'], values = [ledstat])
            self.updateHeaderInfo(header_flag='ext', keys = ['LEDT01','LEDT02','LEDT03','LEDT04','LEDT05','LEDT06','LEDT07','LEDT08','LEDT09','LEDT10','LEDT11','LEDT12','LEDT13','LEDT14'], values = letts)
        

    if obs_param["shutter_effect"] == True:
        pf_map = pf_map * chip.shutter_img
        pf_map = np.array(pf_map, dtype='float32')
        self.updateHeaderInfo(header_flag='ext', keys = ['SHTSTAT'], values = [True])
    else:
        self.updateHeaderInfo(header_flag='ext', keys = ['SHTSTAT','SHTOPEN0','SHTOPEN1','SHTCLOS0','SHTCLOS1'], values = [True,'','','',''])

    chip.img = chip.img + pf_map

    # renew header info
    datetime_obs = datetime.utcfromtimestamp(pointing.timestamp)
    datetime_obs = datetime_obs.replace(tzinfo=timezone.utc)
    t_obs = Time(datetime_obs)
    
    ##ccd刷新2s,等待0.5s，开灯后等待0.5s,开始曝光
    t_obs_renew = Time(t_obs.mjd - (2.) / 86400., format="mjd")

    t_obs_utc = datetime.utcfromtimestamp(np.round(datetime.utcfromtimestamp(t_obs_renew.unix).replace(tzinfo=timezone.utc).timestamp(), 1))
    self.updateHeaderInfo(header_flag='prim', keys = ['DATE-OBS'], values = [t_obs_utc.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-5]])

    #dark time : 
    self.updateHeaderInfo(header_flag='ext', keys = ['DARKTIME'], values = [pointing.exp_time])

    gc.collect()
    return chip, filt, tel, pointing