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

package stats

import (
	"math"

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

// Windowed (rolling) reductions over a series: each output element
// summarises one window of consecutive samples. The result holds
// n − window + 1 elements, one per full window, aligned so element i
// summarises samples [i, i+window). Series of every real dtype are
// accepted, and the extrema follow the package's NaN rule: a NaN never
// wins a comparison, and a window that holds nothing but NaN answers
// NaN, exactly as core.Min and core.Max do.

func rollingCheck(a *core.Array, window int) ([]float64, int, error) {
	const name = "Rolling"
	if a.NDim() != 1 {
		return nil, 0, base.Errf("%s: the series must be a vector, got shape %s", name, base.ShapeText(a.Shape()))
	}
	if a.Dtype() == core.Complex {
		return nil, 0, base.Errf("%s: complex series have no ordering to reduce", name)
	}
	n := a.Len()
	if window < 1 || window > n {
		return nil, 0, base.Errf("%s: the window must lie in [1, %d], got %d", name, n, window)
	}
	// Promote through FloatAt, never through RawFloats alone: the
	// float64 payload is empty for an int or float32 array, and a
	// strided view would be read at the wrong stride. A dense float64
	// array copies its payload directly, the same elements the
	// accessor walk returned. The copy also means the fold below
	// costs no per-element bounds check.
	src := make([]float64, n)
	if fs := rawFloats(a); fs != nil {
		copy(src, fs)
	} else {
		for i := range n {
			src[i] = a.FloatAt(i)
		}
	}
	return src, n - window + 1, nil
}

// RollingMean averages each window of the series.
func RollingMean(a *core.Array, window int) (*core.Array, error) {
	src, outLen, err := rollingCheck(a, window)
	if err != nil {
		return nil, err
	}
	out := core.New(core.Float, outLen)
	rollingTotals(src, out.RawFloats(), outLen, window, true)
	return out, nil
}

// RollingSum totals each window of the series.
func RollingSum(a *core.Array, window int) (*core.Array, error) {
	src, outLen, err := rollingCheck(a, window)
	if err != nil {
		return nil, err
	}
	out := core.New(core.Float, outLen)
	rollingTotals(src, out.RawFloats(), outLen, window, false)
	return out, nil
}

// rollingTotals fills the first outLen entries of vals with each full
// window's total of src: the sum, or the mean for mean.
//
// Past a short window the total moves through the series instead of
// rescanning it: it is carried between positions as a two-float pair,
// so one position costs the entering element plus the negated leaving
// one, whatever the window's length, where the rescan costs a fold
// over the whole window at every position. The pair, not a bare
// running sum, is what keeps that honest: a subtractive update sheds
// the low bits of every add and subtract, and on a long window of
// mixed magnitudes the drift ends up wider than what a per-window
// rescan loses, while the pair carries every bit the format holds.
//
// Short windows stay with the rescan, and that is a measured choice,
// not a concession: the rescan's fold costs a couple of cycles a
// sample while the carried walk pays its compensation chain at every
// position whatever the window, so below the crossover the rescan is
// the faster walk by two to three times, and equally accurate.
//
// A window that answers a non-finite total, through a non-finite
// sample or an overflowed sum, answers exactly what the per-window
// rescan answered: the carried pair turns non-finite with it, the
// window is folded from scratch, and the refolded pair becomes the
// carried state for the windows that follow.
func rollingTotals(src []float64, vals []float64, outLen, window int, mean bool) {
	if window < rollingIncrementalWindow {
		for i := range outLen {
			total := 0.0
			for _, v := range src[i : i+window] {
				total += v
			}
			if mean {
				total /= float64(window)
			}
			vals[i] = total
		}
		return
	}
	hi, lo := 0.0, 0.0
	for _, v := range src[:window] {
		hi, lo = rollingAdd(hi, lo, v)
	}
	total := hi + lo
	if !rollingFinite(total) {
		for _, v := range src[:window] {
			total += v
		}
	}
	if mean {
		total /= float64(window)
	}
	vals[0] = total
	for i := 1; i < outLen; i++ {
		hi, lo = rollingAdd(hi, lo, src[i+window-1])
		hi, lo = rollingAdd(hi, lo, -src[i-1])
		total = hi + lo
		if !rollingFinite(total) {
			// The window answers what the per-window rescan answered,
			// and the carried state restarts from the window itself,
			// refolded with the same two-float care the walk carries:
			// seeding the state from the rescan's rounded total would
			// leak that rounding into every window after.
			hi, lo, total = 0, 0, 0
			for _, v := range src[i : i+window] {
				total += v
				hi, lo = rollingAdd(hi, lo, v)
			}
		}
		if mean {
			total /= float64(window)
		}
		vals[i] = total
	}
}

// rollingIncrementalWindow is the window length the rolling totals
// switch from the per-window rescan to the carried update at. The
// rescan's cost per position grows with the window while the carried
// walk's is flat, and the crossover sits near a window of forty;
// sixty-four is the power of two above it, where the carried walk
// already answers twice as fast.
const rollingIncrementalWindow = 64

// rollingAdd returns the two-float pair for hi + lo + b. Knuth's
// TwoSum catches the rounding error of the wide add, and the
// renormalisation folds the pair back so the low word stays at the
// rounding level of the high one; the pair then represents the carried
// total to double-double precision across an unbounded walk of adds
// and subtracts.
func rollingAdd(hi, lo, b float64) (float64, float64) {
	s := hi + b
	bb := s - hi
	lo += (hi - (s - bb)) + (b - bb)
	t := s + lo
	return t, lo - (t - s)
}

// rollingFinite reports whether v is a finite number: the carried
// total is trusted only while it stays one.
func rollingFinite(v float64) bool {
	return math.Abs(v) <= math.MaxFloat64
}

// RollingMax tracks each window's largest sample. A NaN never wins a
// comparison and an all-NaN window answers NaN, the package's rule.
func RollingMax(a *core.Array, window int) (*core.Array, error) {
	src, outLen, err := rollingCheck(a, window)
	if err != nil {
		return nil, err
	}
	return rollingExtreme(src, outLen, window, true), nil
}

// RollingMin tracks each window's smallest sample, with the same NaN
// rule as RollingMax.
func RollingMin(a *core.Array, window int) (*core.Array, error) {
	src, outLen, err := rollingCheck(a, window)
	if err != nil {
		return nil, err
	}
	return rollingExtreme(src, outLen, window, false), nil
}

// rollingExtreme folds each window with a monotonic deque: an index
// leaves the back of the deque only when a later sample is strictly
// better, so the front always holds the window's extreme, and each
// index enters and leaves the deque once. The scan is O(n) where a
// rescan of every window is O(n·window), and it selects the sample the
// rescan selected: a comparison is strict in both, so a tie, the ±0
// pair included, keeps the earlier index, and a window whose samples
// are all NaN answers its last element, exactly the value the rescan's
// own seed walk ends on.
func rollingExtreme(src []float64, outLen, window int, greater bool) *core.Array {
	out := core.New(core.Float, outLen)
	vals := out.RawFloats()
	// The deque holds indices in ascending order, improving towards the
	// back; head is its front, and every entry before head has expired.
	// The buffer is compacted once the dead prefix outgrows the live
	// region, which keeps it proportional to the deque's depth rather
	// than to the series length: each compaction copies at most the
	// entries it drops, and every entry is dropped once.
	deque := make([]int, 0, min(window, rollingDequeCompact))
	head := 0
	for i, v := range src {
		if head >= rollingDequeCompact && head >= len(deque)-head {
			deque = deque[:copy(deque, deque[head:])]
			head = 0
		}
		if !math.IsNaN(v) {
			if greater {
				for len(deque) > head && src[deque[len(deque)-1]] < v {
					deque = deque[:len(deque)-1]
				}
			} else {
				for len(deque) > head && src[deque[len(deque)-1]] > v {
					deque = deque[:len(deque)-1]
				}
			}
			deque = append(deque, i)
		}
		if i+1 < window {
			continue
		}
		oldest := i - window + 1
		for head < len(deque) && deque[head] < oldest {
			head++
		}
		if head == len(deque) {
			// Every sample of the window is NaN; the rescan answers the
			// window's last element, and so does this.
			vals[oldest] = v
			continue
		}
		vals[oldest] = src[deque[head]]
	}
	return out
}

// rollingDequeCompact is the dead-prefix length at which the windowed
// extrema compact their deque buffer: eight kilobytes of indices, past
// which the copy pays for itself on any series long enough to reach it.
const rollingDequeCompact = 1 << 10