main.py

import argparse, sys, tqdm, time, os, yaml
from glob import glob
from datetime import datetime
import traceback

import numpy as np
from .target import spectrum_generator, target_file_load
from .optics import focal_mask, focal_convolve
from .camera import CosmicRayFrameMaker, sky_frame_maker, CpicVisEmccd
from .io import save_fits, log
from .config import update_able_keys, relative_time
from .config import config as default_config


def vis_observation(
        target: dict,
        skybg: float,
        expt: float,
        nframe: int,
        band: str,
        emset: int,
        obsid: int = 41000000000,
        rotation: float = 0,
        shift: list = [0, 0],
        gnc_info: dict = {},
        csst_format: bool = True,
        camera = CpicVisEmccd(),
        crmaker = CosmicRayFrameMaker(),
        nsample: int = 1,
        emgain=None,
        prograss_bar=None,
        output='./'
) -> np.ndarray:
    """Observation simulation main process on visable band using EMCCD.

    Parameters
    -----------
    target : dict
        target dictionary. See the input of `target.spectrum_generator` for more details.
    skybg : float
        sky background in unit of magnitude/arcsec^2
    expt: float
        exposure time in unit of second
    nframe: int
        number of frames
    band: str
        band name
    emset: int
        EM gain setting value. 1023(0x3FF) for ~1.0× EM gain.
    obsid: int
        observation ID. Start from 4 for CPIC, 01 for science observation. See the input of io.obsid_parser for more details.
    rotation: float
        rotation of the telescope. in unit of degree. 0 means North is up.
    shift: list
        target shift in unit of arcsec
    gnc_info: dict
        gnc_info dictionary. See the input of io.primary_hdu for more details.
    csst_format: bool
        if True, the output fits file will be in the csst format.
    crmaker: CosmicRayFrameMaker
        CosmicRayFrameMaker object. See the input of camera.CosmicRayFrameMaker for more details.
    nsample: int
        number of samples for wide bandpass.
    emgain: float or None
        if None, emgain are set using emset parameter. Else, emgain are set using emgain parameter.
    prograss_bar: bool
        if True, a prograss_bar will be shown.
    output: str
        the output directory.

    Returns
    -------
    image_cube: np.ndarray
    """
    start_time = time.time()
    platescale = default_config['platescale']
    iwa = default_config['mask_width'] / 2
    area = default_config['aperature_area']
    expt_start = camera.system_time

    target_list = []
    if 'cstar' in target.keys():
        target_list = spectrum_generator(target)

    image_cube = []

    if emgain is None:
        emgain_value = camera.emgain_set(emset, self_update=False)
    else:
        emgain_value, emset = emgain, -emgain

    params = {
        'target': target,
        'skybg': skybg,
        'expt': expt,
        'nframe': nframe,
        'band': band,
        'emset': emset,
        'emgain': emgain_value,
        'obsid': obsid,
        'rotation': rotation,
        'shift': shift,
    }

    all_frame_info = []

    if prograss_bar:
        pg_bar = tqdm.tqdm(total=nframe, leave=False)

    for i in range(nframe):
        if prograss_bar:
            pg_bar.update(1)
        else:
            print(f'\r Simulation Running: Frame {i+1}/{nframe}', end='')

        frame_info = {}
        frame_info['expt_start'] = camera.system_time
        focal_frame = focal_convolve(
            band,
            target_list,
            init_shifts=shift,
            rotation=rotation,
            nsample=nsample,
            platesize=camera.flat_shape
        )

        if skybg is None or skybg > 100:
            sky_bkg_frame = 0
        else:
            sky_bkg_frame = sky_frame_maker(
                band,
                skybg,
                platescale,
                camera.flat_shape
            )
            sky_bkg_frame = sky_bkg_frame * area * expt
            sky_bkg_frame = focal_mask(sky_bkg_frame, iwa, platescale)

        cr_frame = crmaker.make_cr_frame(camera.dark_shape, expt)

        camera.time_syn(expt, readout=True)
        image = camera.readout(
            focal_frame + sky_bkg_frame,
            emset,
            expt,
            image_cosmic_ray=cr_frame,
            emgain=emgain
        )
        image_cube.append(image)
        frame_info['expt_end'] = camera.system_time
        frame_info['chiptemp'] = camera.ccd_temp
        frame_info['platescale'] = platescale
        frame_info['iwa'] = iwa
        all_frame_info.append(frame_info)

    image_cube = np.array(image_cube)

    expt_end = camera.system_time
    utc0 = default_config['utc0']
    utc0_stamp = datetime.timestamp(datetime.fromisoformat(utc0))
    expt_start_iso = datetime.fromtimestamp(utc0_stamp + expt_start)
    expt_end_iso = datetime.fromtimestamp(utc0_stamp + expt_end)

    params['EXPSTART'] = expt_start_iso.isoformat()
    params['EXPEND'] = expt_end_iso.isoformat()
    params['frame_info'] = all_frame_info

    save_fits(image_cube, params, gnc_info, camera.__dict__.copy(), csst_format=csst_format, output_folder=output)
    if prograss_bar:
        pg_bar.close()
        print(f'  Done [{time.time() - start_time:.1f}s]                             ')
    else:
        print(f'\r  Done [{time.time() - start_time:.1f}s]                             ')
    return image_cube

def quick_run_v2(
        target_str: str,
        band: str,
        expt: float,
        emgain: float,
        nframe: int,
        skybg: float = None,
        rotation: float = 0,
        shift: list = [0, 0],
        emset_input: bool=False,
        cr_frame: bool=True,
        camera_effect: bool=True,
        prograss_bar=False,
        output='./') -> np.ndarray:
    
    """A quick run function for CPIC.

    Parameters
    ----------
    target_str: str
        target string. See the input of `target.target_file_load` for more details.
    band: str
        band name
    expt: float
        exposure time in unit of second
    emgain: float
        EM gain value. Note that emgain is not the same as emset, controled by emset_input keyword.
    nframe: int
        number of frames
    skybg: float or None
        sky background in unit of magnitude. If None, sky background will not be set.
    rotation: float
        rotation of the telescope. in unit of degree. 0 means North is up.
    shift: list
        target shift in unit of arcsec
    emset_input: bool
        if True, emgain paramter is emset. Else, emgain is set using emgain parameter.
    cr_frame: bool
        if True, cosmic ray frame will be added to the image.
    camera_effect: bool
        if True, camera effect will be added to the image.
    prograss_bar: bool
        if True, a prograss_bar will be shown.
    output: str
        the output directory.

    Returns
    -------
    image_cube: np.ndarray
    """
    print(f'Quick Run: {target_str}')


    log.debug(
        f"""input parameters:
        target_str: {target_str}
        skybg: {skybg}
        band: {band}
        expt: {expt}
        nframe: {nframe}
        emgain: {emgain}
        rotation: {rotation}
        shift: {shift}
        emset_input: {emset_input}
        cr_frame: {cr_frame}
        camera_effect: {camera_effect}
        prograss_bar: {prograss_bar}
        output: {output}
""")

    target_dict = target_file_load(target_str)
    if target_dict == {}:
        target_dict['name'] = 'blank'

    emgain_value = emgain
    if emset_input:
        emgain_value = None

    camera = CpicVisEmccd()
    if not camera_effect:
        camera.switch['bias_vp'] = False
        camera.switch['bias_hp'] = False
        camera.switch['bias_ci'] = False
        camera.switch['bias_shift'] = False 
        camera.switch['badcolumn'] = False
        camera.switch['flat'] = False
        camera.switch['cte'] = False
        camera.switch['shutter'] = False
        camera.switch['nonlinear'] = False
    if not cr_frame:
        camera.switch['cosmicray'] = False

    return vis_observation(
        target=target_dict,
        skybg=skybg,
        expt=expt,
        nframe=nframe,
        band=band,
        emset=emgain,
        emgain=emgain_value,
        csst_format=False,
        shift=shift,
        rotation=rotation,
        camera=camera,
        prograss_bar=prograss_bar,
        output=output
    )



def deduplicate_names_add_count(names: list):
    """remove duplicate names and add count"""
    for i in range(len(names)-1,-1,-1):
        if names.count(names[i]) > 1:
            names[i] = names[i] + '_' + str(names.count(names[i]))


def observation_simulation_from_config(obs_file, config_file):
    """Run observation simulation from config file

    Parameters
    -----------
    obs_file: str
        observation file or folder.
    config_file: str
        config file.

    Examples:
        see examples in `example` folder. 
    """
    config = {}
    if config_file:
        with open(config_file) as fid:
            config = yaml.load(fid, Loader=yaml.FullLoader)

    camera_config = config.get('camera', [{}])
    all_camera = []
    all_camera_name = []
    for c_config in camera_config:
        all_camera.append(CpicVisEmccd(c_config))
        all_camera_name.append(c_config.get('name', 'camera'))
    
    deduplicate_names_add_count(all_camera_name)

    for key in config.keys():
        if key in update_able_keys:
            default_config[key] = config[key]   

    output_folder = default_config['output']
    csst_format = default_config['csst_format']
    nsample = default_config['nsample']
    
    obs_file = os.path.abspath(obs_file)
    file_list = []

    if os.path.isdir(obs_file):
        file_list = glob(f"{obs_file}/*.yaml")
    elif '.yaml' == obs_file[-5:]:
        file_list = [obs_file]
    
    if not file_list:
        log.warning(f"No observation file found in {obs_file}")


    for ind_target, file in enumerate(file_list):
        try:
            with open(file, 'r') as fid:
                obs_info = yaml.load(fid, Loader=yaml.FullLoader)

            target = target_file_load(obs_info.get('target', {}))
            skybg = obs_info.get('skybg', None)
            expt = obs_info['expt']
            band = obs_info['band']
            emset = obs_info['emset']
            nframe = obs_info['nframe']
            obsid = obs_info['obsid']
            rotation = obs_info.get('rotation', 0)
            shift = obs_info.get('shift', [0, 0])
            gnc_info = obs_info.get('gnc_info', {})
            time = obs_info.get('time', 0)
            emgain = obs_info.get('emgain', None)
            time = relative_time(time)

        except Exception as e:
            log.error(f"{file} is not a valid yaml file.")
            log.error(f"Failed with {type(e).__name__}{e}.\n\n {traceback.format_exc()}")
            continue

        ind_camera = 0
        for camera_name, camera in zip(all_camera_name, all_camera):
            ind_camera += 1
            ind_run = ind_target * len(all_camera) + ind_camera
            all_run = len(all_camera) * len(file_list)
            info_text = f"({ind_run}/{all_run}) obsid[{obsid}] with {camera_name}"

            log.info(info_text)
            if time == 0:
                camera.time_syn(time, initial=True)
            else:
                dt = time - camera.system_time
                if dt < 0:
                    log.warning(f'Time is not synced. {dt} seconds are added.')
                    dt = 0
                camera.time_syn(dt, readout=False)

            if len(all_camera) > 1:
                output = os.path.join(output_folder, camera_name)
            else:
                output = output_folder

            try:
                vis_observation(
                    target,
                    skybg,
                    expt,
                    nframe,
                    band,
                    emset,
                    obsid,
                    emgain=emgain,
                    rotation=rotation,
                    shift=shift,
                    gnc_info=gnc_info,
                    camera=camera,
                    output=output,
                    nsample=nsample,
                    csst_format=csst_format,
                    prograss_bar=True)
            except Exception as e:
                log.error(f"{info_text} failed with {type(e).__name__}{e}.\n\n {traceback.format_exc()}")

def main(argv=None):
    """
    Command line interface of csst_cpic_sim
    
    Parameters
    -----------
    argv: list
        input arguments. Default is None. If None, sys.argv is used.
    
    """
    parser = argparse.ArgumentParser(description='Cpic obsevation image simulation')
    parser.set_defaults(func=lambda _: parser.print_usage())

    subparsers = parser.add_subparsers(help='type of runs')
    parser_quickrun = subparsers.add_parser('quickrun', help='a quick observation with no configration file')
    parser_quickrun.add_argument('target_string', type=str, help='example: *5.1/25.3(1.3,1.5)/22.1(2.3,-4.5)')
    parser_quickrun.add_argument('expt', type=float, help='exposure time [ms]')
    parser_quickrun.add_argument('emgain', type=float, help='emgain or emgain set value if emgain_input is False')
    parser_quickrun.add_argument('nframe', type=int, help='number of frames')
    parser_quickrun.add_argument('-b', '--band', type=str, default='f661', help='band, one of f565/f661/f743/f883')
    parser_quickrun.add_argument('-r', '--rotation', type=float, default=0, help='rotation angle [degree]')
    parser_quickrun.add_argument('-s', '--skybk', type=float, default=21, help='magnitude of sky background [mag/arcsec^2]')
    parser_quickrun.add_argument('-f', '--cr_frame', action='store_true', help='if True, cosmic ray frame will be added')
    parser_quickrun.add_argument('-e', '--emset', action='store_true', help='if True, emgain set value will be used as input')
    parser_quickrun.add_argument('-c', '--camera_effect', action='store_true', help='if True, camera effect will be added')
    parser_quickrun.add_argument('-o', '--output', type=str, default='./', help='output folder')

    def quick_run_call(args):
        quick_run_v2(
            target_str=args.target_string,
            expt=args.expt,
            emgain=args.emgain,
            nframe=args.nframe,
            band=args.band,
            rotation=args.rotation,
            skybg=args.skybk,
            cr_frame=args.cr_frame,
            emset_input=args.emset,
            camera_effect=args.camera_effect,
            output=os.path.abspath(args.output),
            prograss_bar=True
        )

    parser_quickrun.set_defaults(func=quick_run_call)

    parser_run = subparsers.add_parser('run', help='observation with configration files or folder')
    parser_run.add_argument('target_config', type=str, help='configration file or folder')
    parser_run.add_argument('-c', '--config', type=str, help='instrument configration file')

    def run_all(args):
        observation_simulation_from_config(
            args.target_config,
            args.config,
        )

    parser_run.set_defaults(func=run_all)

    args = parser.parse_args(argv)
    args.func(args)
    

# if __name__ == '__main__':  # pragma: no cover
#     sys.exit(main())


#     target_example = {
#         'cstar': {
#             'magnitude': 1,
#             'ra': '120d',
#             'dec': '40d',
#             'distance': 10,
#             'sptype': 'F0III',
#         },
#         'stars': [
#             {
#                 'magnitude': 20,
#                 'pangle': 60,
#                 'separation': 1,
#                 'sptype': 'F0III'
#             }
#         ]
#     }
# #     quick_run('', 10, 'f661', 1, 1, 30)
# #     quick_run('*2.4/10(3,5)/15(-4,2)', 21, 'f661', 1, 1, 30)

# #     # normal target
#     observation_simulation(
#         target=target_example,
#         skybg=21,
#         expt=1,
#         nframe=2,
#         band='f661',
#         emgain=30,
#         obsid=51012345678,
#     )

#     # bias
    # observation_simulation(
    #     target=target_example,
    #     skybg=999,
    #     expt=1,
    #     nframe=2,
    #     band='f661',
    #     emgain=1,
    #     obsid=51012345678,
    #     shift=[3, 3],
    #     rotation=60
    # )

#     # bias-gain
#     observation_simulation(
#         target={},
#         skybg=999,
#         expt=0.01,
#         nframe=2,
#         band='f661',
#         emgain=1000,
#         obsid=50012345678,
#     )

#     # dark
#     observation_simulation(
#         target={},
#         skybg=999,
#         expt=100,
#         nframe=2,
#         band='f661',
#         emgain=30,
#         obsid=50112345678,
#     )

#     # flat
#     observation_simulation(
#         target={},
#         skybg=15,
#         expt=10,
#         nframe=2,
#         band='f661',
#         emgain=30,
#         obsid=50112345678,
#     )