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

package stats

import (
	"math"
	"slices"

	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// Robust location and scale summaries: the pair that survives a few
// wild samples the mean and standard deviation would happily chase.

// MedianAbsoluteDeviation returns the median of |x − median(x)|, the
// robust scale estimate that breaks down only when nearly half the
// sample is wild. Multiply by 1.4826 to read it as a standard
// deviation on Gaussian data.
func MedianAbsoluteDeviation(a *core.Array) (float64, error) {
	median, err := Median(a)
	if err != nil {
		return 0, err
	}
	deviations := core.New(core.Float, a.Len())
	vals := deviations.RawFloats()
	// The walk is bounded by the sample's own element count: a rebased
	// view's payload may run past its visible elements, and writing one
	// per payload slot would overrun the deviations buffer.
	if fs := rawFloats(a); fs != nil {
		for i, v := range fs[:a.Len()] {
			vals[i] = math.Abs(v - median)
		}
	} else {
		for i := range a.Len() {
			vals[i] = math.Abs(a.FloatAt(i) - median)
		}
	}
	return Median(deviations)
}

// TrimmedMean averages the sample after dropping the fraction from
// each tail: fraction 0.1 discards the smallest and the largest ten
// percent (floored to whole samples) before averaging, which keeps
// the mean honest against one-sided contamination. fraction must lie
// in [0, 0.5) and leave at least one sample in the middle; a non-finite
// sample is an error.
func TrimmedMean(a *core.Array, fraction float64) (float64, error) {
	if a.Dtype() == core.Complex {
		return 0, base.Errf("TrimmedMean: complex samples have no ordering")
	}
	n := a.Len()
	if n == 0 {
		return 0, base.Errf("TrimmedMean: an empty sample has no mean")
	}
	if !(fraction >= 0) || fraction >= 0.5 {
		return 0, base.Errf("TrimmedMean: the fraction must lie in [0, 0.5), got %g", fraction)
	}
	if err := checkFinite("TrimmedMean", "the sample", a); err != nil {
		return 0, err
	}
	trim := int(fraction * float64(n))
	if n-2*trim < 1 {
		return 0, base.Errf("TrimmedMean: trimming %d samples a side of %d leaves nothing", trim, n)
	}
	sortedVals := make([]float64, n)
	if fs := rawFloats(a); fs != nil {
		copy(sortedVals, fs)
	} else {
		for i := range n {
			sortedVals[i] = a.FloatAt(i)
		}
	}
	slices.Sort(sortedVals)
	window := sortedVals[trim : n-trim]
	// Two-pass scaled summation: dividing through by the largest
	// magnitude first keeps every partial sum inside [−k, k] for a
	// window of k samples, so a mean of huge values cannot overflow
	// into an infinity the way the direct accumulation does. Kahan
	// compensation only repairs rounding, never overflow, so the
	// scaling is what carries the guarantee here.
	mag := 0.0
	for _, v := range window {
		if m := math.Abs(v); m > mag {
			mag = m
		}
	}
	if mag == 0 {
		return 0, nil
	}
	total := 0.0
	for _, v := range window {
		total += v / mag
	}
	// The division comes before the multiplication: the scaled total is
	// at most the window's size in magnitude, so dividing by the count
	// first keeps the final scaling inside the representable range,
	// where multiplying the raw total back by the magnitude could
	// overflow before the division runs.
	return total / float64(n-2*trim) * mag, nil
}