"""
Identifier:     csst-l1/msc/csst_msc_instrument/csst_msc_instrument/image/gain.py
Name:           gain.py
Description:    Gain map related function.
Author:         Li Shao (shaoli@nao.cas.cn)
Created:        2023-12-22
Modified-History:
    2023-12-22, Li Shao, created
    2024-02-06, Li Shao, move to csst_msc_instrument
"""

import numpy as np

from .config import get_array_config

__all__ = ['apply_gain', 'make_gainmap_from_channel_value', ]


def apply_gain(image: np.ndarray,
               gainmap: np.ndarray,
               ) -> np.ndarray:
    """
    Apply gain map.

    The gain map is in unit of electrons/ADU. The output image will be in unit of electrons or electrons/s.

    Parameters
    ----------
    image : numpy.ndarray
        Input image, in unit of ADU or ADU/s.
    gainmap : numpy.ndarray
        Gain map, in unit of electrons/ADU.

    Returns
    -------
    numpy.ndarray
        Output image, in unit of electrons or electrons/s.
    """

    return image * gainmap


def make_gainmap_from_channel_value(gain_values: list[float] | tuple[float] | np.ndarray,
                                    array_config_name: str = 'CCDArrayConfig',
                                    image_format: str = 'raw',
                                    ) -> np.ndarray:
    """
    Make gain map from channel gain values.

    Parameters
    ----------
    gain_values : list[float] or tuple[float] or numpy.ndarray
        Gain value of each channel.
    array_config_name : str, optional
        The name of the pixel array configuration class. Default is "CCDArrayConfig".
    image_format : str, optional
        The output image format: "raw", "cube" or "sky". See config.CCDArrayConfig().

    Returns
    -------
    numpy.ndarray
        Output gain map.
    """

    aconf = get_array_config(array_config_name)(image_format)
    if len(gain_values) != aconf.nchan:
        raise ValueError('The length of input gain values does not match with number of channels.')

    full_shape = aconf.get_shape(full=True)
    output = np.zeros(full_shape, dtype=np.float32)
    for i in range(aconf.nchan):
        cutout = aconf.get_cutout(i+1, output, rtype='channel') + gain_values[i]
        aconf.set_cutout(i+1, output, cutout, rtype='channel')

    return output