// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package stats

import (
	"sourcedock.dev/petrbalvin/tensor/internal/base"
)

import (
	"cmp"
	"math"
	"slices"
)

// Multiple-testing corrections over a vector of p-values: Bonferroni,
// Holm's step-down and Benjamini-Hochberg's step-up, each returning
// the adjusted p-value vector the caller can threshold directly. The
// Holm and Benjamini-Hochberg procedures are defined through a running
// extreme over the sorted p-values, which enforces the monotonicity
// their adjusted values must show: equal or larger raw p-values can
// never receive smaller adjustments, and the enforced running extreme
// maps every sorted position back to its own index. Every adjustment
// is at least the raw p-value it belongs to and never leaves [0, 1].

// Bonferroni returns the Bonferroni-adjusted p-values: each p scaled
// by the vector length m and clamped at 1, the family-wise error rate
// control that asks nothing of the dependence between the tests.
func Bonferroni(p []float64) ([]float64, error) {
	const name = "Bonferroni"
	if err := checkPValues(name, p); err != nil {
		return nil, err
	}
	out := make([]float64, len(p))
	for i, v := range p {
		out[i] = min(1, v*float64(len(p)))
	}
	return out, nil
}

// Holm returns the Holm step-down adjusted p-values. Sorted ascending,
// the ith smallest receives the multiplier m−i and the running maximum
// over its predecessors enforces the non-decreasing order the step-down
// procedure implies; the adjusted vector is then mapped back through
// the original positions.
func Holm(p []float64) ([]float64, error) {
	const name = "Holm"
	if err := checkPValues(name, p); err != nil {
		return nil, err
	}
	m := len(p)
	order := make([]int, m)
	for i := range order {
		order[i] = i
	}
	slices.SortFunc(order, func(a, b int) int { return cmp.Compare(p[a], p[b]) })
	out := make([]float64, m)
	running := 0.0
	for rank, idx := range order {
		running = max(running, float64(m-rank)*p[idx])
		// The max against the raw p is arithmetic beltwork: the
		// multiplier never drops below 1, and this pins the guarantee
		// exactly rather than to rounding.
		out[idx] = min(1, max(running, p[idx]))
	}
	return out, nil
}

// BenjaminiHochberg returns the Benjamini-Hochberg step-up adjusted
// p-values, the q-values of the false discovery rate literature.
// Sorted ascending, the ith smallest receives the multiplier m/(i+1)
// and the running minimum over its successors enforces the
// non-decreasing order the step-up procedure implies; the adjusted
// vector is then mapped back through the original positions.
func BenjaminiHochberg(p []float64) ([]float64, error) {
	const name = "BenjaminiHochberg"
	if err := checkPValues(name, p); err != nil {
		return nil, err
	}
	m := len(p)
	order := make([]int, m)
	for i := range order {
		order[i] = i
	}
	slices.SortFunc(order, func(a, b int) int { return cmp.Compare(p[a], p[b]) })
	out := make([]float64, m)
	running := 1.0
	for i := m - 1; i >= 0; i-- {
		idx := order[i]
		running = min(running, float64(m)/float64(i+1)*p[idx])
		out[idx] = min(1, max(running, p[idx]))
	}
	return out, nil
}

// checkPValues validates a p-value vector for the corrections: it must
// be non-empty, every entry finite, and every entry inside [0, 1], the
// refusals reported with the offending index and value.
func checkPValues(name string, p []float64) error {
	if len(p) == 0 {
		return base.Errf("%s: the p-value vector must not be empty", name)
	}
	for i, v := range p {
		if math.IsNaN(v) || math.IsInf(v, 0) {
			return base.Errf("%s: p[%d] is not finite (%g)", name, i, v)
		}
		if v < 0 || v > 1 {
			return base.Errf("%s: p[%d] = %g lies outside [0, 1]", name, i, v)
		}
	}
	return nil
}