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

package core

// Discrete differences: the workhorse behind finite-difference
// derivatives and signal detrending, along any single axis.

// Diff takes the successive differences along one axis, order times:
// order 1 is out[i] = a[i+1] − a[i] along the axis, order 2 applies
// it again, and so on. The axis shrinks by order; the axis must
// therefore hold more elements than the order, and axis must name one
// of the array's axes. Complex arrays are fine: differences carry no
// ordering assumption. Int arrays keep their dtype, because the
// difference of two int64 is the int64 difference; everything else
// produces float.
func Diff(a *Array, order, axis int) (*Array, error) {
	if order < 1 {
		return nil, errf("Diff: the order must be at least 1, got %d", order)
	}
	if axis < 0 || axis >= a.NDim() {
		return nil, errf("Diff: axis %d is outside the %d axes of shape %s", axis, a.NDim(), shapeText(a.Shape()))
	}
	if a.Shape()[axis] <= order {
		return nil, errf("Diff: axis %d holds %d elements, more than the order %d is needed",
			axis, a.Shape()[axis], order)
	}
	cur := a
	for range order {
		next, err := diffOnce(cur, axis)
		if err != nil {
			return nil, err
		}
		cur = next
	}
	return cur, nil
}

// diffOnce applies one round of differences along the axis. Int
// differences stay int64 and complex ones stay complex; every other
// dtype produces float, the only route that used to widen the int side
// through float64 and round neighbours above 2^53 together.
//
// The output is walked run by run: for one position along the trailing
// dimensions the two neighbours sit a fixed stride apart, so a run of
// the output is a plain elementwise subtraction with the dtype
// dispatched once, not a per-element coordinate fold.
func diffOnce(a *Array, axis int) (*Array, error) {
	if !a.isContiguous() {
		// A strided view's payload window is not the run the walk needs,
		// so reduce it to a dense copy first; the elements, and with
		// them the differences, are the ones the accessors returned.
		a = a.materialise()
	}
	shape := a.Shape()
	outShape := append([]int{}, shape...)
	outShape[axis]--
	dt := Float
	switch a.dt {
	case Complex:
		dt = Complex
	case Int:
		dt = Int
	}
	out := &Array{shape: outShape, dt: dt}
	out.alloc(out.Len())
	tail := 1
	for d := axis + 1; d < len(shape); d++ {
		tail *= shape[d]
	}
	head := 1
	for d := range axis {
		head *= shape[d]
	}
	n := shape[axis]
	switch dt {
	case Int:
		diffRuns(out.ints, a.ints[:a.Len()], head, n, tail)
	case Complex:
		diffRuns(out.complexes, a.complexes[:a.Len()], head, n, tail)
	default:
		// The source widens exactly as FloatAt widens it; a float64
		// source is read in place.
		diffRuns(out.floats, floatPayload(a), head, n, tail)
	}
	return out, nil
}

// diffRuns fills dst with the successive differences of src along an
// axis of n elements that steps by tail elements, for each of head
// outer positions.
func diffRuns[T int64 | float64 | complex128](dst, src []T, head, n, tail int) {
	for h := range head {
		base := h * n * tail
		dstBase := h * (n - 1) * tail
		for i := range tail {
			s, d := base+i, dstBase+i
			for k := range n - 1 {
				dst[d+k*tail] = src[s+(k+1)*tail] - src[s+k*tail]
			}
		}
	}
}