"""
Identifier:     KSC-SJ1-MSC-MBI/stats.py
Name:           stats.py
Description:    some stats programs
Author:         Hu Zou (zouhu@nao.cas.cn)
Created:        2022-11-11
Modified-History:
    2022-11-11, created
    2023-12-04, add docstring
"""

import astropy.stats as ast
import numpy as np


def valid_coordinates(coordinates, size=None):
    """ convert tuple, list or np.ndarray of cooridinates to a 2d ndarray and check the coordinates within an image size
    Parameters
    ----------
        coordinates: 2d array for coordinates
        size: size of an image
    Returns:
    ----------
        coord: reshape coordinates
        indcoord: index of coordinates in a size range
    """
    if isinstance(coordinates, (list, tuple, np.ndarray)):
        coord = np.atleast_2d(coordinates)
        if coord.shape[0] != 2 and coord.shape[1] != 2:
            raise ValueError(
                "coordinates should have at least one axis with 2 elements")
        if coord.shape[1] != 2 and coord.shape[0] == 2:
            coord = coord.transpose()
    else:
        raise TypeError(
            "coordinates should be list or array of (x,y) pixel positions")

    if size is None:
        return coord
    else:
        if len(size) != 2:
            raise ValueError("size should have 2 elements")
        nx, ny = size
        x = coord[:, 0]
        y = coord[:, 1]
        indcoord = np.arange(coord.shape[0])
        good = (x >= 0.5) & (x < nx+0.5) & (y >= 0.5) & (y < ny+0.5)
        if np.any(good):
            indcoord = indcoord[good]
        else:
            raise ValueError('coordinates are not in the image range')
        return coord, indcoord


def closest_match(coord1, coord2, min_dist=1.0):
    """ find closest pairs between two sets of coordinates

    Parameters
    ----------
    coord1: 2d array
        coordinates to be matched
    coord2: 2d array
        coordinates matched to
    min_dist: float
        separation tolerance

    Returns
    -------
    idx1: 1d array
        matched index for coord1
    idx2: 1d array
        matched index for coord2
    """
    coord1 = valid_coordinates(coord1)
    coord2 = valid_coordinates(coord2)

    n1 = len(coord1)
    n2 = len(coord2)
    index1 = []
    index2 = []
    x2 = coord2[:, 0]
    y2 = coord2[:, 1]
    for i in range(n1):
        ix1, iy1 = coord1[i]
        index = np.where((np.abs(x2-ix1) < min_dist) &
                         (np.abs(y2-iy1) < min_dist))[0]
        nmatch = len(index)
        if nmatch < 1:
            continue
        if nmatch > 1:
            x2tmp = x2[index]
            y2tmp = y2[index]
            dist = ((x2tmp-ix1)**2+(y2tmp-iy1)**2)
            indsort = np.argsort(dist)
            index1.append(i)
            index2.append(index[indsort])
        else:
            index1.append(i)
            index2.append(index)
    return index1, index2


def sigmaclip_limitsig(data, error=None, sig_limit=None, **kwd):
    """ sigma-clipping the data with upper limits
    Parameters
        data: 1d array
        error: error corresponding to the data
        sig_limit: the upper limit of the std of the clipped data
        **kwd: keywords for astropy.stats.sigma_clip
    OUPUTS:
        mdata: masked data
    """
    data = np.array(data)
    mdata = ast.sigma_clip(data, **kwd)
    if sig_limit is not None:
        while True:
            med = np.ma.median(mdata)
            sig = np.ma.std(mdata)
            if sig < sig_limit:
                break
            index = np.ma.argmax(np.ma.abs(mdata-med))
            mdata.mask[index] = True
    return mdata


def weighted_mean(x, sigmax, weight_square=True):
    """
    Calculate the mean and estimated errors for a set of data points
    DESCRIPTION:
        This routine is adapted from Program 5-1, XFIT, from "Data Reduction
        and Error Analysis for the Physical Sciences", p. 76, by Philip R.
        Bevington, McGraw Hill.  This routine computes the weighted mean using
        Instrumental weights (w=1/sigma^2).
    Parameters:
        x      - Array of data points
        sigmax - array of standard deviations for data points
        weight_square - if True, weight is invariance, else the reciprocal of the error
    Returns:
        xmean  - weighted mean
        sigmam - standard deviation of mean
        stdm - standard deviation of data
    """
    x = np.atleast_1d(x).copy()
    sigmax = np.atleast_1d(sigmax).copy()
    if len(x) == 1:
        xmean = x
        sigmam = sigmax
        stdm = sigmax
    else:
        weight = 1.0/sigmax**2
        weight1 = weight
        if not weight_square:
            weight1 = 1.0/sigmax
        wsum = weight.sum()
        xmean = (weight1*x).sum()/weight1.sum()
        sigmam = np.sqrt(1.0/wsum)
        stdm = np.sqrt(np.sum((x-xmean)**2)/len(x))
    return xmean, sigmam, stdm
