"""
Identifier:     csst_msc_mbi_instrument/csst_msc_mbi_instrument.py
Name:           csst_msc_mbi_instrument.py
Description:    Make the correction for the intrument effect of CSST MS-MBI
Author:         Zhou Xie, created
Created:        2023-11-27
Modified-History:
    2023-11-27, Zhou Xie, created
"""

from datetime import datetime
from typing import Optional

# import csst_common
import numpy as np
from itertools import product
from astropy.io import fits
from astropy.stats import sigma_clipped_stats
# from csst_common import CsstResult, CsstStatus
# from csst_msc_instrument import (
#     remove_cr_deepcr,
#     subtract_bias,
#     subtract_dark,
#     FlagConfig,
# )
from L1_pipeline.csst_msc_instrument.image import (
    remove_cr_deepcr,
    subtract_bias,
    subtract_dark,
    FlagConfig,
)


SATURATE = 50000


def core_msc_l1_mbi_instrument(
    image_path: str = "/path/to/image",
    bias_path: str = "/path/to/bias",
    dark_path: str = "/path/to/dark",
    flat_path: str = "/path/to/flat",
    shutter_path: str = "/path/to/shutter",
    image_output_path: str = "/path/to/image_output",
    weight_output_path: str = "/path/to/weight_output",
    flag_output_path: str = "/path/to/flag_output",
    deepcr_model_path: Optional[str] = None,
    config_ccd_info: Optional[str] = None,
    config_bad_pixel: Optional[str] = None,
    device: Optional[str] = "CPU",
):  # -> CsstResult:
    """
    Make the instrument correction for one chip of CSST data.

    This function corrects the instrument effect, such as bias, dark, flat, shutter, cosmicrays etc.

    Parameters
    ----------
    image_path : str
        The sci file path.
    bias_path : str
        The bias file path.
    dark_path : str
        The dark file path.
    flat_path : str
        The flat file path.
    shutter_path : str
        The shutter file path.
    image_output_path : str
        The image output file path.
    weight_output_path : str
        The weight output file path.
    flag_output_path : str
        The flag output file path.
    deepcr_model_path : str, optional
        The deepcr model file path.
    config_ccd_info : str, optional
        The config for ccd file path.
    config_bad_pixel : str, optional
        The config for bad pixel file path.
    device : str
        The deepcr device input 'CPU' or 'GPU', by default 'CPU'.

    Returns
    -------
    CsstResult
        Result containing `status` and `files`.
    """
    status_header = _set_status_header()

    # CCD image processing
    with fits.open(image_path) as raw:
        hdu = fits.HDUList(
            [
                fits.PrimaryHDU(header=raw[0].header.copy()),
                fits.ImageHDU(header=raw[1].header.copy(),
                              data=raw[1].data.copy()),
            ]
        )
        # 创建wht[1].data hdulist
        wht = fits.HDUList(
            [
                fits.PrimaryHDU(header=raw[0].header.copy()),
                fits.ImageHDU(
                    header=raw[1].header.copy(),
                    data=np.zeros_like(hdu[1].data, dtype=np.float32),
                ),
            ]
        )
        # 创建flg[1].data hdulist
        flg = fits.HDUList(
            [
                fits.PrimaryHDU(header=raw[0].header.copy()),
                fits.ImageHDU(
                    header=raw[1].header.copy(),
                    data=np.zeros_like(hdu[1].data, dtype=np.uint32),
                ),
            ]
        )

        img_err = np.zeros_like(hdu[1].data)
        # # shutter
        # shutter = fits.getdata(shutter_path, 1)
        bias, bias_error = _read_ccds(bias_path)
        dark, dark_error = _read_ccds(dark_path)
        flat, flat_error = _read_ccds(flat_path)

        hdu[1].data, img_err = subtract_bias(
            hdu[1].data,
            img_err,
            bias,
            bias_error,
        )
        status_header.set("S_BIAS", 0)

        hdu[1].data, img_err = subtract_dark(
            hdu[1].data, img_err, dark, dark_error, hdu[0].header["EXPTIME"]
        )
        status_header.set("S_DARK", 0)

        hdu[1].data, img_err = _divide_flat(
            hdu[1].data,
            img_err,
            flat,
            flat_error,
        )
        status_header.set("S_FLAT", 0)
        # 非线性改正
        hdu[1].data = _fix_nonlinear(hdu[1].data)
        status_header.set("S_NLIN", 0)
        # 坏像素检测
        flg[1].data = _check_badpix(flg[1].data, flat)
        # 热像素和暖像素检测
        flg[1].data = _check_hot_and_warm_pix(
            flg[1].data, dark, hdu[0].header["EXPTIME"], hdu[1].header["RON01"]
        )
        # 过饱和检测
        saturate = SATURATE
        flg[1].data = _check_over_saturation(
            flg[1].data, hdu[1].data, saturate)
        saturate = saturate * \
            hdu[1].header["GAIN01"] / hdu[0].header["EXPTIME"]
        status_header.set("SATURATE", saturate)
        # 时间归一
        hdu[1].data = hdu[1].data / hdu[0].header["EXPTIME"]
        # 单位转换
        hdu[1].data = _multiply_gain_map(hdu[1].data, hdu[1].header)
        # hdu[1].header.set('BUNIT', 'electron/s')
        status_header.set("BUNIT", "electron/s")
        # 检测宇宙线
        hdu[1].data, flg[1].data, cr_count = remove_cr_deepcr(
            hdu[1].data,
            flg[1].data,
            device,
            hdu[1].header["CHIPID"],
        )
        status_header.set("S_CRS", 0)
        status_header.set("CRCOUNT", cr_count)

        # 记录BKG与RMS
        _, bkg, rms = sigma_clipped_stats(
            data=hdu[1].data, mask=(flg[1].data == 16))
        status_header.set("SKY_BKG", bkg)
        status_header.set("SKY_RMS", rms)
        # 检测卫星拖尾
        flg[1].data = _check_satellitetrail(flg[1].data, hdu[1].data)
        status_header.set("S_SAT", 0)
        # 生成wht[1].data
        wht[1].data = _create_weight(
            hdu[1].data, flg[1].data, hdu[1].header["RON01"], hdu[0].header["EXPTIME"]
        )
        # 数据类型
        hdu[1].data = hdu[1].data.astype(np.float32)

        # # 快门改正
        # hdu[1].data = _fix_shutter(hdu[1].data, shutter, hdu[0].header["EXPTIME"])
        # status_header.set("S_SHUT", 0)

        if all(
            [
                status_header["S_BIAS"] == 0,
                status_header["S_DARK"] == 0,
                status_header["S_FLAT"] == 0,
                status_header["S_CRS"] == 0,
                status_header["S_NLIN"] == 0,
                # status_header["S_CTE"] == 0,  # 这个功能还未添加
                status_header["S_SAT"] == 0,
                # status_header["S_SHUT"] == 0,
                status_header["SKY_BKG"] != -9999,
                status_header["SKY_RMS"] != -9999,
                status_header["SATURATE"] != -9999,
                status_header["CRCOUNT"] != -9999,
            ]
        ):
            status_header.set("S_INST", 0)

        # hdu[1].header.extend(status_header, bottom=True)
        # 修正wht和flg数据类型
        wht[1].header = hdu[1].header.copy()
        wht[1].data = wht[1].data.astype(np.float32)
        wht[1].header.remove("BUNIT")

        flg[1].header = hdu[1].header.copy()
        flg[1].data = flg[1].data.astype(np.uint32)
        flg[1].header.set("BUNIT", "dimensionless unit")

        hdu.writeto(image_output_path, overwrite=True)
        wht.writeto(weight_output_path, overwrite=True)
        flg.writeto(flag_output_path, overwrite=True)

    # # construct CsstResult
    # result = CsstResult(
    #     status=CsstStatus(status_header["S_INST"]),  # 0/1/2\
    #     header=status_header,
    #     files=[
    #         image_output_path,
    #         weight_output_path,
    #         flag_output_path,
    #     ],
    # )
    # return result


def _set_status_header():
    header_path = __file__.replace(
        "csst_msc_mbi_instrument.py", "instrument_cor_demo.head"
    )
    demo_header = fits.getheader(header_path, ignore_missing_simple=True)

    status_header = fits.Header()  # template Header
    status_header.append(("COMMENT", "=" * 66), bottom=True)
    status_header.append(
        ("COMMENT", "INSTRUMENT CORRECTION INFORMATION"), bottom=True)
    status_header.append(("COMMENT", "=" * 66), bottom=True)

    for card in demo_header.cards:
        status_header.append(card, bottom=True)

    return status_header


def _read_ccds(
    ref_path: str,
) -> tuple[np.array, np.array]:
    ref = fits.getdata(ref_path)
    return ref, np.zeros_like(ref)


def _divide_flat(
    image: np.ndarray,
    image_error: np.ndarray,
    flat: np.ndarray,
    flat_error: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    def divide(a, b):
        return np.divide(a, b, out=np.zeros_like(a, float), where=(b != 0))

    result = divide(image, flat)
    result_error = (
        np.abs(result)
        * ((divide(image_error, image)) ** 2 + (divide(flat_error, flat)) ** 2) ** 0.5
    )
    return result, result_error


def _fix_nonlinear(image: np.ndarray, beta1: float = 5e-7) -> np.ndarray:
    # 计算非线性系数
    beta = 5e-7
    f1 = [1]
    fnlin0 = 1 + beta
    fnlin = [fnlin0]
    for i in range(1000, 65000, 1000):
        f1.append(i)
        fnlin.append(1 + beta * i)
    # 插值函数
    from scipy.interpolate import PchipInterpolator as PCHIP

    # PCHIP stands for Piecewise Cubic Hermite Interpolating Polynomial
    interp_func = PCHIP(f1, fnlin)
    # 非线性系数矩阵
    imgnlin = interp_func(image)
    # 图像改正
    image = imgnlin * image
    return image


def _check_badpix(
    flag: np.ndarray,
    flat: np.ndarray,
) -> np.ndarray:
    med = np.median(flat)
    _flag = (flat < 0.5 * med) | (1.5 * med < flat)

    flagconf = FlagConfig()
    _flag = _flag.astype(flagconf.dtype)

    value = flagconf.get_value("dead_pixel")
    flag = flag | (_flag * value)
    return flag


def _check_hot_and_warm_pix(
    flag: np.ndarray,
    dark: np.ndarray,
    exptime: float,
    rdnoise: float,
) -> np.ndarray:

    flagconf = FlagConfig()
    _dark = dark * exptime
    _dark[_dark < 0] = 0
    _flag = 1 * rdnoise**2 <= _dark  # 不确定是否包含 暂定包含
    _flag = _flag.astype(flagconf.dtype)
    value = flagconf.get_value("hot_pixel")
    flag = flag | (_flag * value)

    _flag = (0.5 * rdnoise**2 < _dark) & (_dark < 1 * rdnoise**2)
    _flag = _flag.astype(flagconf.dtype)
    value = flagconf.get_value("warm_pixel")
    flag = flag | (_flag * value)
    return flag


def _check_over_saturation(
    flag: np.ndarray,
    image: np.ndarray,
    saturate: int = 65535,
    iterations: int = 0,
) -> np.ndarray:
    _flag = image >= saturate
    _flag = _dilation(_flag, iterations=iterations)
    flagconf = FlagConfig()
    _flag = _flag.astype(flagconf.dtype)
    value = flagconf.get_value("saturated")
    flag = np.bitwise_or(flag, _flag * value)
    return flag


def _dilation(array_orig: np.ndarray, iterations: int = 1):
    """
    Make a dilation for array

    This function makes a dilation for the 2D mask array (saturated pixels)

    Parameters
    ----------
    array_orig : np.ndarray
        The mask array without dilation.
    structure : int
        The number of dilations performed on the structure with itself.

    Returns
    -------
    array_out : np.ndarray
        The mask array after dilation.
    """
    from scipy import ndimage

    struct1 = ndimage.generate_binary_structure(
        2, 1)  # give a saturate structure
    struct_ext = ndimage.iterate_structure(
        struct1, iterations)  # make a dilation
    if iterations == 0:
        array_out = array_orig
    else:
        array_out = ndimage.binary_dilation(
            array_orig, structure=struct_ext, iterations=1
        ).astype(array_orig.dtype)

    return array_out


def _check_satellitetrail(
    flag: np.ndarray,
    image: np.ndarray,
) -> np.ndarray:
    from skimage import exposure, transform
    from skimage.draw import line
    from skimage.feature import canny

    # 调整图像
    p1, p2 = np.percentile(image, (0.1, 99.9))
    image = exposure.rescale_intensity(image, in_range=(p1, p2))
    # 边界识别，转化为二值图像
    immax = np.max(image)
    # canny边界
    edge = canny(
        image=image, sigma=2.0, low_threshold=immax * 0.1, high_threshold=immax * 0.5
    )

    # 概率霍夫变换
    angle = np.radians(np.arange(2, 178, 0.5, dtype=float))
    # 直线端点
    result = transform.probabilistic_hough_line(
        image=edge, threshold=210, line_length=400, line_gap=75, theta=angle
    )
    result = np.asarray(result)
    # 绘制mask
    _flag = np.zeros(image.shape)
    # 遍历每一条线段
    for (x1, y1), (x2, y2) in result:
        xx, yy = line(x1, y1, x2, y2)
        _flag[yy, xx] = 1
    flagconf = FlagConfig()
    _flag = _flag.astype(flagconf.dtype)
    value = flagconf.get_value("satellite_trail")
    flag = np.bitwise_or(flag, _flag * value)
    return flag


def _create_weight(
    img: np.ndarray,
    flg: np.ndarray,
    rdnoise: float,
    exptime: float,
) -> np.ndarray:
    data = img.copy()
    data[img < 0] = 0
    var_raw = data * exptime + rdnoise**2
    var_bias = 0.0
    weight = 1.0 / (var_raw + var_bias) * exptime**2
    weight[flg > 0] = 0
    wht = weight
    return wht


def _fix_shutter(
    img: np.ndarray, shutter: np.ndarray, exptime: float = 150.0
) -> np.ndarray:
    img_cor = np.float32(img * exptime * 1000.0 / (exptime * 1000.0 + shutter))
    return img_cor


def _multiply_gain_map(
    image: np.ndarray,
    header: fits.header,
) -> np.ndarray:
    gain_map = np.zeros_like(image)
    h = header["NAXIS1"] // 8
    w = header["NAXIS2"] // 2
    for y, x in product(range(8), (0, 1)):
        y = y if x == 0 else 7 - y
        gain = header["GAIN" + str(x * 8 + y + 1).zfill(2)]
        gain_map[slice(w * x, w * (x + 1)), slice(h * y, h * (y + 1))] = gain
        # gain_map[w * x : w * (x + 1), h * y : h * (y + 1)] = gain
    return image * gain_map


# if __name__ == "__main__":
#     result: CsstResult = core_msc_l1_mbi_instrument(
#         image_path=r"D:\Zhou\Desktop\data\data_09.fits",
#         bias_path=r"D:\Zhou\Desktop\data\bias_09.fits",
#         dark_path=r"D:\Zhou\Desktop\data\dark_09.fits",
#         flat_path=r"D:\Zhou\Desktop\data\flat_09.fits",
#         shutter_path=r"D:\Zhou\Desktop\data\csst_msc_ms_shutter_09_000001.fits",
#         image_output_path=r"D:\Zhou\Desktop\data\result\image_test_output.fits",
#         weight_output_path=r"D:\Zhou\Desktop\data\result\weight_test_output.fits",
#         flag_output_path=r"D:\Zhou\Desktop\data\result\flag_test_output.fits",
#     )
#     print(result.status == 0)
