"""
Statistical computations and bootstrap resampling routines (converted from stats_subs.f90).
"""

import numpy as np


def get_samplevec(nsamp, rng=None):
    """Generate a random resampling vector (bootstrap notation).

    The initial data set corresponds to a sampling vector of [1, 1, 1, ...],
    whereas each resampling vector has the same sum (resampling with replacement),
    but with different integer weights on each entry.

    Parameters
    ----------
    nsamp : int
        Number of samples in input data vector.
    rng : numpy.random.Generator, optional
        Random number generator. If None, uses default.

    Returns
    -------
    sampv : ndarray of int, shape (nsamp,)
        Output sampling vector.
    """
    if rng is None:
        rng = np.random.default_rng()

    sampv = np.zeros(nsamp, dtype=np.int32)
    indices = rng.integers(0, nsamp, size=nsamp)
    for idx in indices:
        sampv[idx] += 1
    return sampv


def resample_data(data_in, sampv):
    """Apply a resampling vector to an input data vector.

    Parameters
    ----------
    data_in : ndarray, shape (nsamp,)
        Input data vector.
    sampv : ndarray of int, shape (nsamp,)
        Sampling vector.

    Returns
    -------
    data_out : ndarray, shape (nsamp,)
        Resampled data vector.
    """
    data_out = np.repeat(data_in, sampv)
    return data_out


def mean(x):
    """Compute the mean of array x."""
    if len(x) == 0:
        return 0.0
    return np.mean(x)


def median(x):
    """Compute the median of array x."""
    if len(x) == 0:
        return 0.0
    return np.median(x)


def mean_stddev(x):
    """Compute mean and sample standard deviation of array x.

    Returns
    -------
    xmean : float
    xstd : float (sample std, N-1 normalization)
    """
    n = len(x)
    if n == 0:
        return 0.0, 0.0
    xmean = np.mean(x)
    if n == 1:
        return xmean, 0.0
    xstd = np.std(x, ddof=1)
    return xmean, xstd


def piksrt(x):
    """Sort array x in place (insertion sort). Returns sorted copy."""
    return np.sort(x)


def robomean2(x, xgap, nit):
    """Robust mean (L1/L2 hybrid) using IRLS-type iterative weighting.

    Parameters
    ----------
    x : ndarray
        Input data vector.
    xgap : float
        Distance/weighting parameter.
    nit : int
        Number of iterations.

    Returns
    -------
    xmean : float
        Weighted mean.
    fit2 : float
        Robust L2 norm.
    """
    n = len(x)
    if n == 0:
        return 0.0, 0.0

    xw = np.ones(n)

    for it in range(nit):
        sumw = np.sum(xw)
        xmean = np.sum(x * xw) / sumw

        if it < nit - 1:
            d = xgap + np.abs(x - xmean)
            xw = 1.0 / d

    fit2 = np.sum(xw * (x - xmean) ** 2)
    return xmean, fit2


def indexx(arr):
    """Return indices that sort array from low to high.

    Parameters
    ----------
    arr : ndarray

    Returns
    -------
    indx : ndarray of int
        Sort indices (0-based).
    """
    return np.argsort(arr)


def getstd(x):
    """Compute the population standard deviation of array x."""
    if len(x) == 0:
        return 0.0
    xmn = np.mean(x)
    return np.sqrt(np.mean((x - xmn) ** 2))


def getmad(x):
    """Compute the Median Absolute Deviation of array x."""
    if len(x) == 0:
        return 0.0
    xmed = np.median(x)
    return np.median(np.abs(x - xmed))


def lexsort_strings(inarray):
    """Get sort indices for a character array (lexicographic).

    Parameters
    ----------
    inarray : list of str

    Returns
    -------
    isort : ndarray of int
        Indices to sort array.
    """
    return np.array(sorted(range(len(inarray)), key=lambda i: inarray[i]))
