"""
Identifier:     csst-l1/msc/csst_msc_instrument/csst_msc_instrument/image/crosstalk.py
Name:           crosstalk.py
Description:    Crosstalk correction.
Author:         Tianmeng Zhang (zhangtm@nao.cas.cn)
Created:        2023-11-16
Modified-History:
    2023-11-16, Tianmeng Zhang, created
    2023-11-22, Li Shao, split from common_detector
    2024-02-06, Li Shao, move to csst_msc_instrument
"""

import numpy as np

__all__ = ["remove_crosstalk", ]


def remove_crosstalk(data: np.ndarray,
                     crosstalk_coe: np.ndarray,
                     hdu_num: int = 16,
                     switch: bool = False,
                     ) -> np.ndarray:
    """
    Function to remove the crosstalk between different channel.

    Function to remove the crosstalk between different channel.

    Parameters
    ----------
    data : numpy.ndarray
        The input data array which is needed to remove the crosstalk from each hdu, 3D-array.
    crosstalk_coe : numpy.ndarray
        The coefficients of crosstalk, in size [hdu_num, hdu_num].
    hdu_num : int, optional
        The number of extension for each detector.
    switch : bool, optional
        If True, do the crosstalk correction. If False, return original data. Default is False.

    Returns
    -------
    numpy.ndarray
        The data array after crosstalk correction, 3D-array.
    """

    data_cor = np.copy(data)  # backup original data array

    if switch:
        for i in range(hdu_num):
            for j in range(hdu_num):
                data_cor[i, :, :] = data[i, :, :] + data[j, :, :] * crosstalk_coe[i, j]
    else:
        data_cor = data
        print('No crosstalk correction, return the original array')

    return data_cor
