main.py

import numpy as np
import re

from .target import spectrum_generator
from .optics import make_focus_image, focal_mask, optics_config
from .psf_simulation import simulate_psf
from .camera import EMCCD, CosmicRayFrameMaker, sky_frame_maker
from .io import save_fits, log
from .config import which_focalplane


def psf_function(band, cstar_spectrum, shape, error=0.1):
    cstar = True
    if shape < 300:
        cstar = False
    return simulate_psf(error, band, cstar_spectrum, nsample=1, cstar=cstar)


def observation_simulation(
        target: dict,
        skybg: float,
        expt: float,
        nframe: int,
        band: str,
        emgain: float,
        obsid: int = 51900000000,
        rotation: float = 0,
        shift: list = [0, 0],
        gnc_info: dict = {},
        csst_format: bool = True,
        psf_function: callable = psf_function):
    """
    Simulate the observation. All-In-One function of the package.

    Parameters
    -----------
    target: dict
        The target information. See target.py for details.
    skybg: float
        magnitude of the skybackground at the input b and. (abmag system)
    expt: float
        exposure time in second.
    nframe: int
        number of frames to be simulated.
    band: str
        the band of the observation. (e.g. 'f661')
    emgain: float
        the EM gain of the camera.
    obsid: int
        the observation id. Default is 51900000000.
    rotation: float
        the rotation angle of the target in degree. Default is 0.
    shift: list
        the shift of the target in arcsec. Default is [0, 0].
    gnc_info: dict
        the gnc information. Default is {}. See io.py for details.
    csst_format: bool
        whether to save the fits file in CSST format. Default is True.
    psf_function: callable
        the function to generate the psf. See optics.py for details.

    Returns
    -----------
    np.ndarray of the simulated images with shape (nframe, 1088, 1050).

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

    focal_name = which_focalplane(band)
    this_focal_config = optics_config[focal_name]
    telescope_config = optics_config['telescope']
    area = telescope_config['aperature_area']

    if focal_name == 'vis':
        camera = EMCCD()
    else:
        raise ValueError('Only VIS focal plane is supported.')

    platescale = this_focal_config['platescale']
    iwa = this_focal_config['mask_width'] / 2

    crmaker = CosmicRayFrameMaker()
    images = []

    params = {
        'target': target,
        'skybg': skybg,
        'expt': expt,
        'nframe': nframe,
        'band': band,
        'emgain': emgain,
        'obsid': obsid,
        'rotation': rotation,
        'shift': shift,
    }
    paramstr = ', '.join([f'{k}={v}' for k, v in params.items()])
    log.debug(f"parameters: {paramstr}")

    for i in range(nframe):

        log.info(f'Simulation Running: Frame {i+1}/{nframe}')

        focal_frame = make_focus_image(
            band,
            target_list,
            psf_function,
            platesize=camera.flat_shape,
            rotation=rotation,
            init_shifts=shift,
        )

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

        focal_frame = (focal_frame + sky_bkg_frame) * area
        focal_frame = focal_mask(focal_frame, iwa, platescale)

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

        img = camera.readout(
            focal_frame,
            emgain,
            expt,
            image_cosmic_ray=cr_frame
        )

        images.append(img)

    images = np.array(images)

    save_fits(images, params, gnc_info, csst_format=csst_format)

    return images


def quick_run(
        target_str: str,
        skymag: float,
        band: str,
        expt: float,
        nframe: int,
        emgain: float,
        rotation: float = 0,
        shift: list = [0, 0]) -> np.ndarray:
    """
    A quick run function to simulate the observation.

    Parameters
    -----------
    target_str: str
        The target information in string format.
        In the format of "\*5.1/25.3(1.3,1.5)/22.1(2.3,-4.5)" which means a central star
        with magnitude 5.1, and two substellar with magnitude 25.3 and 22.1, respectively.
        The first number in the parenthesis is the x position in arcsec, and the second is the y position.
    skybg: float
        magnitude of the skybackground at the input band. (abmag system)
    band: str
        the band of the observation. (e.g. 'f661')
    expt: float
        exposure time in second.
    nframe: int
        number of frames to be simulated.
    emgain: float
        the EM gain of the camera.
    rotation: float (optional)
        the rotation angle of the target in degree. Default is 0.
    shift: list (optional)
        the shift of the target in arcsec. Default is [0, 0].

    Returns
    -----------
    np.ndarray of the simulated images, with shape (nframe, 1088, 1050)

    Notes
    -----------
    1. stars are assumed to be G0III with distance 10pc.
    2. magnitude of the star and substellar are assumed to be in the same band.
    3. Csst format is not supported.
    4. The psf is assumed to be the default one.
    5. fits file will be saved in the current directory.


    """
    log.info(f'Quick Run: {target_str}')
    target_dict = {
        'name': 'cal',
    }

    if (target_str != '') and (target_str[0] == '*'):
        objects = target_str[1:].split('/')
        cstar_mag = float(objects[0])

        cstar = {
            'magnitude': cstar_mag,
            'ra': '0d',
            'dec': '0d',
            'sptype': 'G0III',
            'distance': 10,
            'mag_input_band': band
        }

        stars = []
        for sub_stellar in objects[1:]:

            float_regex = R"[+-]?\d+(?:\.\d+)?"
            match = re.match(
                rf"({float_regex})\(({float_regex}),({float_regex})\)", sub_stellar)
            if not match:
                raise ValueError('Wrong format for sub stellar.')
            mag = float(match.group(1))
            x = float(match.group(2))
            y = float(match.group(3))
            pangle = np.arctan2(x, y) * 180 / np.pi
            separation = np.sqrt(x**2 + y**2)
            stars.append({
                'magnitude': mag,
                'pangle': pangle,
                'separation': separation,
                'sptype': 'G0III',
                'mag_input_band': band
            })
        target_dict = {
            'name': target_str[1:],
            'cstar': cstar,
            'stars': stars,
        }

    return observation_simulation(
        target=target_dict,
        skybg=skymag,
        expt=expt,
        nframe=nframe,
        band=band,
        emgain=emgain,
        csst_format=False,
        shift=shift,
        rotation=rotation,
    )


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

# quick_run('*5.1/25.3(0.8,0.8)', None, 'f661', 10, 1, 10)
# quick_run('*5/20(0.8,0.8)', None, 'f883', 10, 1, 10)

# # quick *5.1/25.3(1.3,1.5) expt nframe emgain band rotation shift
# # quick target_name expt nframe emgain band rotation shift
# # plan plan_file_or_folder
if __name__ == '__main__':  # pragma: no cover
    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,
#     )