"""
Identifier:     csst-l1/msc/csst_msc_common/csst_msc_common/cosmicray/deepcr.py
Name:           deepcr.py
Description:    Cosmic ray correction with DeepCR.
Author:         Tianmeng Zhang (zhangtm@nao.cas.cn), Xie Zhou (xiezhou@cnlab.net)
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
"""

import os
import numpy as np
import deepCR

from ..config import FlagConfig

__all__ = ["remove_cr_deepcr", ]

MODEL_DIR = os.path.join(os.path.dirname(
    os.path.abspath(__file__)), 'cr_model')

DEEPCR_MODEL_PATH = {
    '01': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '02': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '03': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '04': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '05': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '06': 'CSST-CCD06-2023-06-08_hidden32_epoch50.pth',
    '07': 'CSST-CCD07-2023-06-08_hidden32_epoch50.pth',
    '08': 'CSST-CCD08-2023-06-08_hidden32_epoch50.pth',
    '09': 'CSST-CCD09-2023-06-08_hidden32_epoch50.pth',
    '10': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '11': 'CSST-CCD11-2023-06-07_hidden32_epoch50.pth',
    '12': 'CSST-CCD12-2023-06-07_hidden32_epoch50.pth',
    '13': 'CSST-CCD13-2023-06-07_hidden32_epoch50.pth',
    '14': 'CSST-CCD14-2023-06-07_hidden32_epoch50.pth',
    '15': 'CSST-CCD15-2023-06-06_hidden32_epoch50.pth',
    '16': 'CSST-CCD16-2023-06-06_hidden32_epoch50.pth',
    '17': 'CSST-CCD17-2023-06-06_hidden32_epoch50.pth',
    '18': 'CSST-CCD18-2023-06-05_hidden32_epoch50.pth',
    '19': 'CSST-CCD19-2023-06-04_hidden32_epoch50.pth',
    '20': 'CSST-CCD20-2023-06-05_hidden32_epoch50.pth',
    '21': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '22': 'CSST-CCD22-2023-06-05_hidden32_epoch50.pth',
    '23': 'CSST-CCD23-2023-04-27_hidden32_epoch50.pth',
    '24': 'CSST-CCD24-2023-06-05_hidden32_epoch50.pth',
    '25': 'CSST-CCD25-2023-06-05_hidden32_epoch50.pth',
    '26': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '27': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '28': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '29': '2022-04-26_15_samples_32_layers_epoch30.pth',
    '30': '2022-04-26_15_samples_32_layers_epoch30.pth', }


def remove_cr_deepcr(image: np.ndarray,
                     flag: np.ndarray,
                     device: str,
                     chipid: str,
                     ) -> tuple[np.ndarray, np.ndarray, int]:
    """
    Remove cosmic rays.

    Remove cosmic rays.

    Parameters
    ----------
    image : numpy.ndarray
        Image to be processed.
    flag : numpy.ndarray
        Input bitflag map.
    device : str
        Processing unit: "CPU" or "GPU".
    chipid : str
        CCD chip id, like "06" or "25", header['CHIPID'].

    Returns
    -------
    image : numpy.ndarray
        Processed image.
    flag : numpy.ndarray
        Added bitflag map.
    cr_count : int
        Number of flagged cosmic rays.
    """
    inpaint = False
    inpaint_model = 'ACS-WFC-F606W-2-32'
    model_path = os.path.join(MODEL_DIR, DEEPCR_MODEL_PATH[chipid])
    model = deepCR.deepCR(
        model_path,
        inpaint_model,
        device=device.upper(),
        hidden=32
    )
    print("Start running deepCR...")
    results = model.clean(
        image,
        threshold=0.5,
        inpaint=False,
        binary=True,
        segment=True,
        patch=256,
        parallel=False,
        n_jobs=1
    )
    # results = model.clean(
    #     image,
    #     threshold=0.5,
    #     inpaint=False,
    #     binary=True,
    #     segment=True,
    #     patch=1024,
    #     parallel=False,
    #     n_jobs=-1
    # )
    if inpaint:
        _flag, image = results
    else:
        _flag = results
    flagconf = FlagConfig()
    _flag = _flag.astype(flagconf.dtype)
    value = flagconf.get_value('cosmic_ray')
    flag = np.bitwise_or(flag, _flag * value)
    cr_count = int(np.sum(_flag))

    return image, flag, cr_count
