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

package signal

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

import "slices"

import "sourcedock.dev/petrbalvin/tensor/internal/engine"

// Numerically robust scalar reduction and finite-difference stencils
// for grid data. SumKahan applies Kahan's compensated summation to the
// whole array: the same sum `Sum` computes, but with the running
// rounding error carried alongside, so long chains of similar
// magnitudes stop bleeding digits. The finite-difference stencils
// produce the central-difference gradient and Laplacian of grid data,
// the raw material of explicit PDE steppers.

// SumKahan returns the sum of all elements by Kahan's compensated
// summation. For arrays with mixed magnitudes the result carries more
// of the small terms than a plain left-to-right accumulation.
// core.Complex arrays are refused.
func SumKahan(a *core.Array) (float64, error) {
	if a.Dtype() == core.Complex {
		return 0, base.Errf("SumKahan: complex arrays are not supported")
	}
	src := widenFloats(a)
	var sum, comp float64
	// Bounded by the element count: a rebased view's aliased payload
	// runs past its extent.
	for i := range a.Len() {
		v := src[i]
		v -= comp
		t := sum + v
		comp = (t - sum) - v
		sum = t
	}
	return sum, nil
}

// Gradient1D returns the central-difference derivative of a 1-D
// uniform grid signal y with spacing dx. Interior points use the
// second-order central stencil; the endpoints fall back to the
// first-order one-sided stencil, so the result has the same length as
// the input.
func Gradient1D(y *core.Array, dx float64) (*core.Array, error) {
	if y.NDim() != 1 {
		return nil, base.Errf("Gradient1D: needs a rank-1 signal, got shape %s", base.ShapeText(y.Shape()))
	}
	if y.Dtype() == core.Complex {
		return nil, base.Errf("Gradient1D: complex arrays are not supported")
	}
	n := y.Len()
	if n < 2 {
		return nil, base.Errf("Gradient1D: needs at least 2 points, got %d", n)
	}
	if dx == 0 {
		return nil, base.Errf("Gradient1D: spacing must not be zero")
	}
	src := widenFloats(y)
	out := core.New(core.Float, n)
	dst := out.RawFloats()
	dst[0] = (src[1] - src[0]) / dx
	for i := 1; i < n-1; i++ {
		dst[i] = (src[i+1] - src[i-1]) / (2 * dx)
	}
	dst[n-1] = (src[n-1] - src[n-2]) / dx
	return out, nil
}

// Laplacian returns the second-derivative Laplacian of a grid signal.
// 1-D input takes one spacing, 2-D input two (dx, dy for a row-major
// (rows, cols) grid, 5-point stencil), 3-D input three (7-point
// stencil). Boundary points copy their nearest interior value; the
// caller owns the boundary condition. A rank outside 1..3, a zero
// spacing, a missing spacing, or a degenerate extent is an error:
// every axis needs at least 2 points, and every axis of a 2-D or 3-D
// grid at least 3, so the interior stencil and the boundary copies
// have somewhere to stand.
func Laplacian(a *core.Array, spacings ...float64) (*core.Array, error) {
	const name = "Laplacian"
	if a.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex arrays are not supported", name)
	}
	if len(spacings) < a.NDim() {
		return nil, base.Errf("%s: needs %d spacings for rank %d, got %d",
			name, a.NDim(), a.NDim(), len(spacings))
	}
	if slices.Contains(spacings, 0) {
		return nil, base.Errf("%s: spacing must not be zero", name)
	}
	switch a.NDim() {
	case 1:
		if a.Len() < 2 {
			return nil, base.Errf("%s: needs at least 2 points, got %d", name, a.Len())
		}
	case 2, 3:
		for i, d := range a.Shape() {
			if d < 3 {
				return nil, base.Errf("%s: axis %d has extent %d, the stencil needs at least 3 points",
					name, i, d)
			}
		}
	default:
		return nil, base.Errf("%s: needs a 1-D, 2-D or 3-D array, got rank %d", name, a.NDim())
	}
	switch a.NDim() {
	case 1:
		return laplacian1D(a, spacings[0]), nil
	case 2:
		return laplacian2D(a, spacings[0], spacings[1]), nil
	default:
		return laplacian3D(a, spacings[0], spacings[1], spacings[2]), nil
	}
}

// laplacian1D applies the 3-point second-difference stencil along the
// single axis. Interior: (y[i+1] − 2y[i] + y[i−1]) / dx². For n = 2
// there is no interior point, and the edge copy leaves both entries
// zero.
func laplacian1D(y *core.Array, dx float64) *core.Array {
	n := y.Len()
	out := core.New(core.Float, n)
	inv := 1 / (dx * dx)
	src := widenFloats(y)
	dst := out.RawFloats()
	for i := 1; i < n-1; i++ {
		dst[i] = (src[i+1] - 2*src[i] + src[i-1]) * inv
	}
	dst[0] = dst[1]
	dst[n-1] = dst[n-2]
	return out
}

// laplacian2D applies the 5-point stencil on a row-major grid of
// shape (ny, nx) with spacings dx, dy.
func laplacian2D(y *core.Array, dx, dy float64) *core.Array {
	ny, nx := y.Shape()[0], y.Shape()[1]
	out := core.New(core.Float, ny, nx)
	invx := 1 / (dx * dx)
	invy := 1 / (dy * dy)
	src := widenFloats(y)
	dst := out.RawFloats()
	engine.ParallelMin(ny, workFloorFor(nx*6), func(s, e int) {
		for j := s; j < e; j++ {
			for i := 1; i < nx-1; i++ {
				idx := j*nx + i
				c := src[idx]
				acc := (src[idx+1] + src[idx-1] - 2*c) * invx
				if j > 0 && j < ny-1 {
					acc += (src[(j-1)*nx+i] + src[(j+1)*nx+i] - 2*c) * invy
				}
				dst[idx] = acc
			}
			dst[j*nx] = dst[j*nx+1]
			dst[j*nx+nx-1] = dst[j*nx+nx-2]
		}
	})
	for i := range nx {
		dst[i] = dst[nx+i]
		dst[(ny-1)*nx+i] = dst[(ny-2)*nx+i]
	}
	return out
}

// laplacian3D applies the 7-point stencil on a row-major grid of
// shape (nz, ny, nx) with spacings dx, dy, dz. Boundary planes copy
// their nearest interior plane, the caller owns the boundary
// condition.
func laplacian3D(y *core.Array, dx, dy, dz float64) *core.Array {
	nz, ny, nx := y.Shape()[0], y.Shape()[1], y.Shape()[2]
	out := core.New(core.Float, nz, ny, nx)
	invx := 1 / (dx * dx)
	invy := 1 / (dy * dy)
	invz := 1 / (dz * dz)
	src := widenFloats(y)
	dst := out.RawFloats()
	engine.ParallelMin(nz, workFloorFor(ny*nx*9), func(zs, ze int) {
		for k := zs; k < ze; k++ {
			for j := range ny {
				for i := 1; i < nx-1; i++ {
					idx := (k*ny+j)*nx + i
					c := src[idx]
					acc := (src[idx+1] + src[idx-1] - 2*c) * invx
					if j > 0 && j < ny-1 {
						acc += (src[(k*ny+j-1)*nx+i] + src[(k*ny+j+1)*nx+i] - 2*c) * invy
					}
					if k > 0 && k < nz-1 {
						acc += (src[((k-1)*ny+j)*nx+i] + src[((k+1)*ny+j)*nx+i] - 2*c) * invz
					}
					dst[idx] = acc
				}
				// The x boundaries copy their nearest interior value.
				dst[(k*ny+j)*nx] = dst[(k*ny+j)*nx+1]
				dst[(k*ny+j)*nx+nx-1] = dst[(k*ny+j)*nx+nx-2]
			}
		}
	})
	// The y boundaries copy their nearest interior row.
	for k := range nz {
		for i := range nx {
			dst[k*ny*nx+i] = dst[(k*ny+1)*nx+i]
			dst[((k+1)*ny-1)*nx+i] = dst[((k+1)*ny-2)*nx+i]
		}
	}
	// The z boundaries copy their nearest interior plane.
	for j := range ny {
		for i := range nx {
			dst[j*nx+i] = dst[(ny+j)*nx+i]
			dst[((nz-1)*ny+j)*nx+i] = dst[((nz-2)*ny+j)*nx+i]
		}
	}
	return out
}