import numpy as np
from astropy.io import fits


def array_combine(ndarray, mode="mean") -> np.ndarray:
    """ Function to combine 3-D data array

    Parameters
    ----------
    ndarray: array, input data cube (3D)
    mode: mean, median, sum, mean_clip, median_clip, default is mean
    """
    if mode == "median":
        array = np.median(ndarray, axis=0)
    elif mode == "median_clip":
        ndarray = np.sort(ndarray, axis=0)[1:-1]
        array = np.median(ndarray, axis=0)
    elif mode == "sum":
        array = np.sum(ndarray, axis=0)
    elif mode == "mean":
        array = np.mean(ndarray, axis=0)
    elif mode == "mean_clip":
        ndarray = np.sort(ndarray, axis=0)[1:-1]
        array = np.mean(ndarray, axis=0)
    return array


def load_bias(path: str) -> np.ndarray:
    with fits.open(path) as hdul:
        du = hdul[1].data
    du = du.astype(int)
    return du


def load_dark(path: str, bias) -> np.ndarray:
    with fits.open(path) as hdul:
        du = hdul[1].data
        hu = hdul[0].header
    du = du.astype(int)
    du = du - bias
    du = du / hu["EXPTIME"]
    return du


def load_flat(path: str, bias, dark) -> np.ndarray:
    with fits.open(path) as hdul:
        du = hdul[1].data
        hu = hdul[0].header
    du = du.astype(int)
    du = du - bias - dark * hu["EXPTIME"]
    du = du / hu["EXPTIME"]
    du = du / np.median(du)
    return du


def combine(func, mode: str, path_list, *args) -> np.ndarray:
    du_list = [func(path, *args) for path in path_list]
    du = array_combine(du_list, mode)
    return du


def combine_images(b_p_lst, d_p_lst, f_p_lst, mode_list=["median", "median", "median", ]):
    """

    Parameters
    ----------
    b_p_lst:
        List of currently ccd number bias file path
    d_p_lst:
        List of currently ccd number dark file path
    f_p_lst:
        List of currently ccd number flat file path
    mode_list:
        [0] bias combine mode
        [1] dark combine mode
        [2] flat combine mode
        mean, median, sum, mean_clip, median_clip
    """
    bias = combine(load_bias, mode_list[0], b_p_lst)
    dark = combine(load_dark, mode_list[1], d_p_lst, bias)
    flat = combine(load_flat, mode_list[2], f_p_lst, bias, dark)
    return bias.copy(), dark.copy(), flat.copy()
