"""
Identifier:     csst-l1/msc/csst_msc_common/csst_msc_common/cosmicray/lacosmic.py
Name:           lacosmic.py
Description:    Cosmic ray correction with L.A.Cosmic.
Author:         Xie Zhou (xiezhou@cnlab.net), Tianmeng Zhang (zhangtm@nao.cas.cn)
Created:        2023-11-16
Modified-History:
    2023-11-16, Tianmeng Zhang, created
    2023-11-17, Xie Zhou, add lacosmic
    2023-11-22, Li Shao, split from common_cr, improve code format
"""

import numpy as np
from ccdproc import cosmicray_lacosmic

from ..config import FlagConfig

__all__ = ["remove_cr_la", ]


def remove_cr_la(image: np.ndarray,
                 flag: np.ndarray,
                 gain: float = 1.1,
                 readnoise: float = 5.0,
                 sigclip: float = 3.0,
                 sigfrac: float = 0.5,
                 objlim: float = 5.0,
                 satlevel: float = 65535.0,
                 pssl: float = 0.0,
                 niter: int = 4,
                 sepmed: bool = True,
                 cleantype: str = 'meanmask',
                 fsmode: str = 'median',
                 psfmodel: str = 'gauss',
                 psffwhm: float = 2.5,
                 psfsize: int = 7,
                 psfk: np.ndarray | None = None,
                 psfbeta: float = 4.765,
                 verbose: bool = False,
                 gain_apply: bool = True,
                 ) -> tuple[np.ndarray, np.ndarray, int]:
    """
    Remove cosmic rays with L.A.Cosmic algorithm.

    Remove cosmic rays with L.A.Cosmic algorithm.

    Parameters
    ----------
    image : numpy.ndarray
        Image to be processed.
    flag : numpy.ndarray
        Input bitflag map.
    gain : float or astropy.units.Quantity, optional
        Gain of the image (electrons / ADU). Default: 1.0.
    readnoise : float, optional
        Read noise of the image (electrons). Used to generate the noise model of the image. Default: 6.5.
    sigclip : float, optional
        Laplacian-to-noise limit for cosmic ray detection.
        Lower values will flag more pixels as cosmic rays. Default: 4.5.
    sigfrac : float, optional
        Fractional detection limit for neighboring pixels. For cosmic ray neighbor pixels,
        a Laplacian-to-noise detection limit of sigfrac * sigclip will be used. Default: 0.3.
    objlim : float, optional
        Minimum contrast between Laplacian image and the fine structure image. Increase this value if
        cores of bright stars are flagged as cosmic rays. Default: 5.0.
    satlevel : float, optional
        Saturation level of the image (electrons). This value is used to detect saturated stars and pixels
        at or above this level are added to the mask. Default: 65535.0.
    pssl : float, optional
        Previously subtracted sky level in ADU. Default: 0.0.
    niter : int, optional
        Number of iterations of the LA Cosmic algorithm to perform. Default: 4.
    sepmed : bool, optional
        Use the separable median filter instead of the full median filter. The separable median is not
        identical to the full median filter, but they are approximately the same, the separable median
        filter is significantly faster, and still detects cosmic rays well. Note, this is a performance
        feature, and not part of the original L.A. Cosmic. Default: True.
    cleantype : str, optional
        Set which clean algorithm is used:
        - "median": An unmasked 5x5 median filter.
        - "medmask": A masked 5x5 median filter.
        - "meanmask": A masked 5x5 mean filter.
        - "idw": A masked 5x5 inverse distance weighted interpolation.
        Default: "meanmask".
    fsmode : str, optional
        Method to build the fine structure image:
        - "median": Use the median filter in the standard LA Cosmic algorithm.
        - "convolve": Convolve the image with the psf kernel to calculate the fine structure image.
        Default: "median".
    psfmodel : str, optional
        Model to use to generate the psf kernel if fsmode == 'convolve' and psfk is None. The current
        choices are Gaussian and Moffat profiles:
        - "gauss" and "moffat" produce circular PSF kernels.
        - The "gaussx"`` and "gaussy" produce Gaussian kernels in the x and y directions respectively.
        Default: "gauss".
    psffwhm : float, optional
        Full Width Half Maximum of the PSF to use to generate the kernel. Default: 2.5.
    psfsize : int, optional
        Size of the kernel to calculate. Returned kernel will have size psfsize x psfsize. psfsize should be odd.
        Default: 7.
    psfk : numpy.ndarray[float] or None, optional
        PSF kernel array to use for the fine structure image if fsmode == "convolve".
        If None and fsmode == "convolve", calculate the psf kernel using "psfmodel". Default: None.
    psfbeta : float, optional
        Moffat beta parameter. Only used if fsmode=="convolve" and psfmodel=="moffat". Default: 4.765.
    verbose : bool, optional
        Print to the screen or not. Default: False.
    gain_apply : bool, optional
        If True, return gain-corrected data, with correct units, otherwise do not gain-correct the data.
        Default is "True" to preserve backwards compatibility.

    Returns
    -------
    image : numpy.ndarray
        Processed image.
    flag : numpy.ndarray
        Added bitflag map.
    cr_count : int
        Number of flagged cosmic rays.
    """

    image, _flag = cosmicray_lacosmic(ccd=image,
                                      sigclip=sigclip,
                                      sigfrac=sigfrac,
                                      objlim=objlim,
                                      gain=gain,
                                      readnoise=readnoise,
                                      satlevel=satlevel,
                                      pssl=pssl,
                                      niter=niter,
                                      sepmed=sepmed,
                                      cleantype=cleantype,
                                      fsmode=fsmode,
                                      psfmodel=psfmodel,
                                      psffwhm=psffwhm,
                                      psfsize=psfsize,
                                      psfk=psfk,
                                      psfbeta=psfbeta,
                                      verbose=verbose,
                                      gain_apply=gain_apply)

    flagconf = FlagConfig()
    _flag = _flag.astype(flagconf.dtype)
    value = flagconf.get_value('cosmic_ray')
    flag = np.bitwise_or(flag, _flag * value)
    cr_count = np.sum(_flag)
    return image, flag, cr_count
