/*  Helper functions
 *
 *  Copyleft (C) 2025-2026, Marek Gagolewski <https://www.gagolewski.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License
 *  Version 3, 19 November 2007, published by the Free Software Foundation.
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 *  GNU Affero General Public License Version 3 for more details.
 *  You should have received a copy of the License along with this program.
 *  If this is not the case, refer to <https://www.gnu.org/licenses/>.
 */


#ifndef __c_argfuns_h
#define __c_argfuns_h


#include "c_common.h"
#include <stdexcept>
#include <algorithm>


/*! Comparer for argsort().
 *
 *  Ensures the resulting permutation is stable.
 */
template<class T>
struct __argsort_comparer
{
    const T* x;

    __argsort_comparer(const T* x) { this->x = x; }

    inline bool operator()(const Py_ssize_t i, const Py_ssize_t j) const {
        return this->x[i] <  this->x[j] ||
              (this->x[i] == this->x[j] && i < j);
    }
};


/*! Finds an(*) ordering permutation w.r.t. `<`.
 *
 *  Both ret and x should be of the same length n;
 *  ret will be overwritten.
 *
 *  (*) or THE stable one, if stable=true, which is the default.
 *
 *  We call permutation o stable, whenever i < j and x[i]==x[j]
 *  implies that o[i] < o[j].
 *
 *  @param ret return array
 *  @param x array to order
 *  @param n size of ret and x
 *  @param stable use a stable sorting algorithm? (slower)
 */
template<class T>
void Cargsort(Py_ssize_t* ret, const T* x, Py_ssize_t n, bool stable=true)
{
    if (n <= 0) return;

    for (Py_ssize_t i=0; i<n; ++i)
        ret[i] = i;

    if (stable)
        std::stable_sort(ret, ret+n, __argsort_comparer<T>(x));
    else
        std::sort(ret, ret+n, __argsort_comparer<T>(x));
}




/*! Returns the index of the (k-1)-th smallest value in an array x.
 *
 *  argkmin(x, 0) == argmin(x), or, more generally,
 *  argkmin(x, k) == np.argsort(x)[k].
 *
 *  Run time: O(nk), where n == len(x). Working mem: O(k).
 *  Does not modify x.
 *
 *  Very fast for small k and randomly ordered
 *  or almost sorted (increasingly) data.
 *
 *
 *  If buf is not NULL, it must be of length at least k+1.
 *
 *  @param x data
 *  @param n length of x
 *  @param k value in {0,...,n-1}, preferably small
 *  @param buf optional working buffer of size >= k+1, will be overwritten
 */
template<class T>
Py_ssize_t Cargkmin(const T* x, Py_ssize_t n, Py_ssize_t k, Py_ssize_t* buf=NULL)
{
    Py_ssize_t* idx;

    if (n <= 0) throw std::domain_error("n <= 0");
    if (k >= n) throw std::domain_error("k >= n");

    k += 1;
    if (!buf) idx = new Py_ssize_t[k];
    else      idx = buf;

    for (Py_ssize_t i=0; i<k; ++i) {
        Py_ssize_t j = i;
        idx[i] = i;
        while (j > 0 && x[i] < x[idx[j-1]]) {
            idx[j] = idx[j-1];
            j -= 1;
        }
        idx[j] = i;
    }

    for (Py_ssize_t i=k; i<n; ++i) {
        if (x[idx[k-1]] <= x[i])
            continue;
        Py_ssize_t j = k-1;
        idx[k-1] = i;
        while (j > 0 && x[i] < x[idx[j-1]]) {
            idx[j] = idx[j-1];
            j -= 1;
        }
        idx[j] = i;
    }

    Py_ssize_t ret = idx[k-1];

    if (!buf) delete [] idx;

    return ret;
}



#endif
