"""
Identifier:     csst-l1/msc/csst_msc_instrument/csst_msc_instrument/image/overscan.py
Name:           overscan.py
Description:    Overscan correction and data format conversion.
Author:         Li Shao (shaoli@nao.cas.cn)
Created:        2023-11-16
Modified-History:
    2023-11-16, Li Shao, created
    2023-11-22, Li Shao, split from common_detector
    2024-02-06, Li Shao, move to csst_msc_instrument
"""

from typing import Callable
import numpy as np
from scipy.ndimage import median_filter
from scipy.signal import savgol_filter
from scipy.special import erf
from scipy.stats import trim_mean, t, sigmaclip
from astropy.stats import mad_std, sigma_clip

from measurement_pipeline.L1_pipeline.csst_msc_instrument.image.config import get_array_config

__all__ = ["average_overscan", "smooth_overscan", "correct_overscan", ]


def average_overscan(image: np.ndarray,
                     axis: int | None = None,
                     cen_method: str | Callable = 'trim_mean',
                     std_method: str | Callable = 'mad_std',
                     clip_flag: bool = False,
                     clip_sigma: float = 3.5,
                     clip_cenfunc: str | Callable = 'median',
                     clip_stdfunc: str | Callable = 'std',
                     ) -> tuple[float | np.ndarray, float | np.ndarray]:
    """
    Average overscan data along one direction or overall.

    Average overscan data along one direction or overall.

    Parameters
    ----------
    image : numpy.ndarray
        Input image: overscan region cutout.
    axis : int, optional
        Axis along which the operation is performed. If None (default), use the full array.
    cen_method : str or Callable, optional
        Algorithm to calculate average value: "trim_mean" (default), "median", "mean" or Callable function.
    std_method : str or Callable, optional
        Algorithm to calculate standard deviation: "mad_std" (default), "std" or Callable function.
    clip_flag : bool, optional
        If True, use sigma-clipping.
    clip_sigma : float, optional
        Clipping threshold for sigma-clipping. Default is 3.5.
    clip_cenfunc : str or Callable, optional
        Center value function for sigma-clipping: "median" (default), "mean" or Callable function.
    clip_stdfunc : str or Callable, optional
        Standard deviation function for sigma-clipping: "std" (default), "mad_std" or Callable function.

    Returns
    -------
    avg : numpy.ndarray
        Averaged overscan data (one value or 1D array).
    err : numpy.ndarray
        The uncertainty.
    """

    # sigma-clipping
    if clip_flag and cen_method != 'trim_mean':
        data = sigma_clip(image, sigma=clip_sigma, axis=axis, masked=True,
                          cenfunc=clip_cenfunc, stdfunc=clip_stdfunc)

        if cen_method == 'median':
            avg = np.ma.median(data)
        elif cen_method == 'mean':
            avg = np.ma.mean(data)
        else:
            avg = cen_method(data)

        if std_method == 'mad_std':
            err = mad_std(data, axis=axis)
        elif std_method == 'std':
            err = np.ma.std(data, axis=axis)
        else:
            err = std_method(data, axis=axis)

    # without clipping
    else:
        if cen_method == 'trim_mean':
            avg = trim_mean(image, 0.1, axis=axis)
        elif cen_method == 'median':
            avg = np.median(image, axis=axis).astype(float)
        elif cen_method == 'mean':
            avg = np.mean(image, axis=axis)
        else:
            avg = cen_method(image, axis=axis)

        if std_method == 'mad_std':
            err = mad_std(image, axis=axis)
        elif std_method == 'std':
            err = np.std(image, axis=axis)
        else:
            err = std_method(image, axis=axis)

    # uncertainty of the average
    ny, nx = image.shape
    if axis is None:
        n = ny * nx
    elif axis == 0:
        n = ny
    else:
        n = nx
    if cen_method == 'trim_mean':
        n = n * 0.8
    err = err / np.sqrt(n) * t.ppf(erf(1), df=n - 1)  # uncertainty of the mean
    if cen_method == 'median' or cen_method is np.median or cen_method is np.ma.median:
        err *= 1.2533

    return avg, err


def smooth_overscan(data: np.ndarray,
                    filter_width: int,
                    filter_method: str = 'savgol',
                    filter_polyorder: int = 2,
                    ) -> np.ndarray:
    """
    Smooth 1D averaged overscan data.

    Smooth 1D averaged overscan data.

    Parameters
    ----------
    data : numpy.ndarray
        Input 1-D overscan data.
    filter_width : int
        The width of filter. Must be positive.
    filter_method : str, optional
        Process algorithm: "savgol" (default) or "median".
    filter_polyorder : int, optional
        The order of polynomial used for Savitzky-Golay filter.

    Returns
    -------
    numpy.ndarray
        Smoothed 1-D overscan data.
    """

    if filter_method == 'savgol':
        dd = savgol_filter(data, filter_width * 2, filter_polyorder)
        diff = data - dd
        index = np.abs(diff) < sigmaclip(diff, 5, 5)[0].std() * 5  # clip outliers
        x = np.arange(len(data))
        dd = np.interp(x, x[index], data[index])
        output = savgol_filter(dd, filter_width, filter_polyorder)
    else:
        output = median_filter(data, size=filter_width)

    return output


def correct_overscan(image: np.ndarray,
                     array_config_name: str = 'CCDArrayConfig',
                     region_type: str = 'x',
                     correct_type: str = 'line',
                     gap: int = 10,
                     cen_method: str | Callable = 'trim_mean',
                     std_method: str | Callable = 'mad_std',
                     clip_flag: bool = False,
                     clip_sigma: float = 3.5,
                     clip_cenfunc: str | Callable = 'median',
                     clip_stdfunc: str | Callable = 'std',
                     filter_width: int = 0,
                     filter_method: str = 'savgol',
                     filter_polyorder: int = 2,
                     clean_inactive: bool = True,
                     ) -> tuple[np.ndarray, np.ndarray]:
    """
    Overscan correction.

    Overscan correction. All parameters are optional inputs.

    Parameters
    ----------
    image : numpy.ndarray
        Input 2-D image.
    array_config_name : str, optional
        The name of the pixel array configuration class. Default is "CCDArrayConfig".
    region_type : str, optional
        The type of overscan region. "x" (default) is serial overscan, "y" is parallel overscan.
    correct_type : str, optional
        Correction type. "line" (default) is correct by row or column. "all" is correct with all pixels.
    gap : int, optional
        Number of pixels after active pixels not to use for overscan correction.
    cen_method : str or Callable, optional
        Algorithm to calculate average value: "trim_mean" (default), "median", "mean" or Callable function.
    std_method : str or Callable, optional
        Algorithm to calculate standard deviation: "mad_std" (default), "std" or Callable function.
    clip_flag : bool, optional
        If True, use sigma-clipping.
    clip_sigma : float, optional
        Clipping threshold for sigma-clipping. Default is 3.5.
    clip_cenfunc : str or Callable, optional
        Center value function for sigma-clipping: "median" (default), "mean" or Callable function.
    clip_stdfunc : str or Callable, optional
        Standard deviation function for sigma-clipping: "std" (default), "mad_std" or Callable function.
    filter_width : int, optional
        The width of filter. If smaller than 2, then smoothing will be skipped.
    filter_method : str, optional
        Filtering algorithm: "savgol" (default) or "median".
    filter_polyorder : int, optional
        The order of polynomial used for Savitzky-Golay filter.
    clean_inactive : bool, optional
        If True, all pixels out of active region are set to zero. Default is True.

    Returns
    -------
    output : numpy.ndarray
        Overscan corrected image.
    output_err : numpy.ndarray
        Uncertainty image.

    Examples
    --------
    image_corr, err_corr = correct_overscan(image, region_type='x', correct_type='line')
    image_corr, err_corr = correct_overscan(image, region_type='y', correct_type='all', clean_inactive=True)
    """

    aconf = get_array_config(array_config_name)('raw')
    if correct_type == 'line':
        if region_type == 'x':
            axis = 1
        else:
            axis = 0
    else:
        axis = None

    if clean_inactive:
        output = np.zeros(image.shape)
    else:
        output = image.astype(float)
    output_err = np.zeros(image.shape)

    for i in range(aconf.nchan):
        overscan = aconf.get_cutout(i+1, image, rtype='{}_overscan'.format(region_type))

        # remove the first pixels (to avoid potential CTI contamination)
        if gap > 0:
            if region_type == 'x':
                overscan = overscan[:, gap:]
            else:
                overscan = overscan[gap:, :]

        # calculate average and uncertainty
        avg_value, err_value = average_overscan(overscan, axis=axis,
                                                cen_method=cen_method, std_method=std_method,
                                                clip_flag=clip_flag, clip_sigma=clip_sigma,
                                                clip_cenfunc=clip_cenfunc, clip_stdfunc=clip_stdfunc)

        # smooth if required
        if filter_width > 1 and axis is not None:
            avg_value = smooth_overscan(avg_value, filter_width,
                                        filter_method=filter_method, filter_polyorder=filter_polyorder)
            err_value = err_value / np.sqrt(filter_width) * 1.2533

        # correct the active region
        x1, x2, y1, y2 = aconf.get_boundary(i+1, rtype='active')
        nx, ny = x2 - x1, y2 - y1
        if axis is None:
            output[y1:y2, x1:x2] = image[y1:y2, x1:x2] - avg_value
            output_err[y1:y2, x1:x2] = err_value
        elif axis == 0:
            output[y1:y2, x1:x2] = image[y1:y2, x1:x2] - np.tile(avg_value, ny).reshape((ny, nx))
            output_err[y1:y2, x1:x2] = np.tile(err_value, ny).reshape((ny, nx))
        else:
            output[y1:y2, x1:x2] = image[y1:y2, x1:x2] - np.repeat(avg_value, nx).reshape((ny, nx))
            output_err[y1:y2, x1:x2] = np.repeat(err_value, nx).reshape((ny, nx))

    return output, output_err
