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

package signal

import (
	"math"

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

// Spectral solution of the Poisson equation on a rectangle with
// homogeneous boundary data. The sine basis vanishes on the boundary
// and diagonalises −Δ for Dirichlet data; the cosine basis has zero
// slope there and does the same for Neumann. Both solves run one
// separable transform pair per axis and one division per mode, and
// both refuse the inputs the mathematics refuses: nonzero-mean
// sources for Neumann, exactly as the periodic solve does.

// poissonRect validates the shared rectangle arguments and returns
// the grid shape and spacings.
func poissonRect(name string, f *core.Array, lx, ly float64) (rows, cols int, hx, hy float64, err error) {
	if f.NDim() != 2 {
		return 0, 0, 0, 0, base.Errf("%s: f must be rank 2, got shape %s", name, base.ShapeText(f.Shape()))
	}
	if f.Dtype() != core.Float {
		return 0, 0, 0, 0, base.Errf("%s: f must be a float64 array, got %s", name, f.Dtype())
	}
	rows, cols = f.Shape()[0], f.Shape()[1]
	if rows < 3 || cols < 3 {
		return 0, 0, 0, 0, base.Errf("%s: the grid must be at least 3×3 to hold interior points, got %d×%d", name, rows, cols)
	}
	if lx <= 0 || ly <= 0 {
		return 0, 0, 0, 0, base.Errf("%s: the side lengths must be positive, got %g and %g", name, lx, ly)
	}
	// A non-finite source would drive the compatibility tolerances NaN
	// and slip through their gates, so it is refused up front. A dense
	// float64 payload is scanned slice-wise: the elements are the ones
	// FloatAt returns.
	vals := poissonFloats(f)
	for _, v := range vals {
		if math.IsInf(v, 0) || math.IsNaN(v) {
			return 0, 0, 0, 0, base.Errf("%s: f holds the non-finite value %g", name, v)
		}
	}
	hx = lx / float64(cols-1)
	hy = ly / float64(rows-1)
	return rows, cols, hx, hy, nil
}

// poissonFloats streams the grid's elements in element order: the raw
// payload when the array is dense, a materialised copy through FloatAt
// when it is a strided view, whose payload order is not the element
// order. Treat an aliased result as read-only.
func poissonFloats(f *core.Array) []float64 {
	if !f.Strided() {
		return f.RawFloats()
	}
	out := make([]float64, f.Len())
	for i := range out {
		out[i] = f.FloatAt(i)
	}
	return out
}

// poissonTrapzSum returns the trapezoidal-weighted sum of f over the
// whole grid: the boundary rows and columns count half, the weight the
// mirrored stencil's own left null vector carries. It is the constant
// mode of the cosine transform pair, the measure the Neumann solve is
// compatible in.
func poissonTrapzSum(f *core.Array) float64 {
	rows, cols := f.Shape()[0], f.Shape()[1]
	vals := poissonFloats(f)
	sum := 0.0
	for r := range rows {
		wr := 1.0
		if r == 0 || r == rows-1 {
			wr = 0.5
		}
		for c := range cols {
			wc := 1.0
			if c == 0 || c == cols-1 {
				wc = 0.5
			}
			sum += wr * wc * vals[r*cols+c]
		}
	}
	return sum
}

// SolvePoissonDirichlet solves −Δu = f on the rectangle with u held
// at zero on the whole boundary. f is sampled on the (rows × cols)
// grid including the boundary rows and columns, whose entries play no
// role in the solve; u comes back on the same grid with a zero
// boundary. Row r of f samples y = r·hy, column c samples x = c·hx,
// with hx = lx/(cols−1) and hy = ly/(rows−1). The sine transform
// pair along each axis diagonalises the interior second-difference
// operator, so the error is the stencil's O(h²), decaying with the
// grid. Non-float dtypes, grids under 3×3 and non-positive side
// lengths are errors.
func SolvePoissonDirichlet(f *core.Array, lx, ly float64) (*core.Array, error) {
	const name = "SolvePoissonDirichlet"
	rows, cols, hx, hy, err := poissonRect(name, f, lx, ly)
	if err != nil {
		return nil, err
	}
	interiorR, interiorC := rows-2, cols-2
	// The interior right-hand side, transformed along both axes with
	// the orthonormal DST-I, whose basis vanishes on the boundary.
	spectrum, err := transformBlock(f, 1, 1, interiorR, interiorC, true)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	// Divide by the stencil eigenvalues: two per axis from the cosine
	// identity on the second difference. Each eigenvalue depends on one
	// index alone, so the per-axis values are built once instead of once
	// per mode; the expressions are the ones the inner loop evaluated.
	eigX := make([]float64, interiorC+1)
	for kx := 1; kx <= interiorC; kx++ {
		eigX[kx] = 2 / (hx * hx) * (1 - math.Cos(math.Pi*float64(kx)/float64(interiorC+1)))
	}
	eigY := make([]float64, interiorR+1)
	for ky := 1; ky <= interiorR; ky++ {
		eigY[ky] = 2 / (hy * hy) * (1 - math.Cos(math.Pi*float64(ky)/float64(interiorR+1)))
	}
	for ky := 1; ky <= interiorR; ky++ {
		ly := eigY[ky]
		row := (ky - 1) * interiorC
		for kx := 1; kx <= interiorC; kx++ {
			spectrum[row+kx-1] /= eigX[kx] + ly
		}
	}
	interior, err := inverseBlock(spectrum, interiorR, interiorC, true)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	out := core.New(core.Float, rows, cols)
	into := out.RawFloats()
	for r := range interiorR {
		for c := range interiorC {
			into[(r+1)*cols+(c+1)] = interior[r*interiorC+c]
		}
	}
	return out, nil
}

// SolvePoissonNeumann solves −Δu = f with zero normal derivative on
// the whole boundary: u comes back on the full grid, boundary points
// included, every sample an unknown. Conventions mirror
// SolvePoissonDirichlet, with the cosine transform pair along each
// axis. The stencil is the ghost-mirror one, the neighbour outside the
// boundary standing in for zero slope, so a boundary row couples its
// mirrored neighbour with 2/h² where an interior row uses 1/h² (the
// boundary control volume is half an interior one). Its cosine modes
// cos(πkx x/lx)·cos(πky y/ly), kx = 0…cols−1 and ky = 0…rows−1, carry
// the eigenvalues 2/hx²·(1−cos(πkx/(cols−1))) +
// 2/hy²·(1−cos(πky/(rows−1))) per axis, and the error against the
// closed form is the stencil's O(h²).
//
// The operator annihilates the constants, so a solution exists only
// for sources whose trapezoidal-weighted sum vanishes; that sum is the
// operator's own compatibility condition (the trapezoidal measure is
// its left null vector), not the plain mean, and a nonzero one is an
// error rather than a quietly shifted problem. The free constant of u
// is fixed by setting the constant mode to zero, so the result has
// zero trapezoidal mean. Non-float dtypes, grids under 3×3 and
// non-positive side lengths are errors.
func SolvePoissonNeumann(f *core.Array, lx, ly float64) (*core.Array, error) {
	const name = "SolvePoissonNeumann"
	rows, cols, hx, hy, err := poissonRect(name, f, lx, ly)
	if err != nil {
		return nil, err
	}
	// The compatibility gate: the constant mode of the transformed
	// source is the trapezoidal-weighted sum, so anything above the
	// accumulated round-off of that sum breaks solvability. The
	// tolerance follows the periodic solve's constant-mode gate, and it
	// is relative to the source: the sum of a trapezoidal grid scales
	// with the sample magnitude and the element count, so an absolute
	// floor would let a tiny source through with a residual far above
	// its own size (a 1e-9 source used to be accepted with a 1.8e-5
	// relative residual).
	sum := poissonTrapzSum(f)
	tol := 64 * base.EpsF * float64(rows*cols) * maxAbsF(f)
	if math.Abs(sum) > tol {
		return nil, base.Errf("%s: f has the nonzero trapezoidal mean %g, the Neumann problem has no solution",
			name, sum/float64((rows-1)*(cols-1)))
	}
	// The pipeline. The mirrored operator M (rows [2,−2] on the
	// boundary, [−1,2,−1] inside, over h²) is diagonalised by the
	// UNNORMALISED DCT-I, whose boundary input weight of 1/2 is the
	// trapezoidal measure; the library carries the orthonormal DCT-I
	// (boundary weight 1/sqrt2), and W = diag(1/2,1,…,1,1/2) turns one
	// into the other. So the forward trip scales by W^(1/2) before the
	// DCT-I, divides by the eigenvalues, and the return trip scales by
	// W^(-1/2) after the second DCT-I, which is its own inverse.
	weighted := core.New(core.Float, rows, cols)
	src := poissonFloats(f)
	dst := weighted.RawFloats()
	for r := range rows {
		sr := 1.0
		if r == 0 || r == rows-1 {
			sr = math.Sqrt2 / 2
		}
		for c := range cols {
			sc := 1.0
			if c == 0 || c == cols-1 {
				sc = math.Sqrt2 / 2
			}
			dst[r*cols+c] = src[r*cols+c] * sr * sc
		}
	}
	spectrum, err := transformBlock(weighted, 0, 0, rows, cols, false)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	// The x eigenvalues depend on kx alone, so the row is built once and
	// read by every ky; the expression is the one the inner loop
	// evaluated. The y eigenvalue was already hoisted to the outer loop.
	eigX := make([]float64, cols)
	for kx := range cols {
		eigX[kx] = 2 / (hx * hx) * (1 - math.Cos(math.Pi*float64(kx)/float64(cols-1)))
	}
	for ky := range rows {
		ly := 2 / (hy * hy) * (1 - math.Cos(math.Pi*float64(ky)/float64(rows-1)))
		row := ky * cols
		for kx := range cols {
			i := row + kx
			if kx == 0 && ky == 0 {
				// The constant mode is the operator's null space: it
				// stays zero, which is what fixes the solution's free
				// constant, instead of dividing by its zero eigenvalue.
				spectrum[i] = 0
				continue
			}
			spectrum[i] /= eigX[kx] + ly
		}
	}
	interior, err := inverseBlock(spectrum, rows, cols, false)
	if err != nil {
		return nil, base.Errf("%s: %w", name, err)
	}
	// The return trip: undo the W^(1/2) scaling, and strip the
	// trapezoidal mean, the measure the solve works in and the one the
	// operator leaves free.
	out := core.New(core.Float, rows, cols)
	into := out.RawFloats()
	sum, weight := 0.0, 0.0
	for r := range rows {
		sr, wr := 1.0, 1.0
		if r == 0 || r == rows-1 {
			sr, wr = math.Sqrt2, 0.5
		}
		for c := range cols {
			sc, wc := 1.0, 1.0
			if c == 0 || c == cols-1 {
				sc, wc = math.Sqrt2, 0.5
			}
			v := interior[r*cols+c] * sr * sc
			into[r*cols+c] = v
			sum += wr * wc * v
			weight += wr * wc
		}
	}
	mean := sum / weight
	for i := range into {
		into[i] -= mean
	}
	return out, nil
}

// maxAbsF returns the largest absolute sample of a float array.
func maxAbsF(f *core.Array) float64 {
	largest := 0.0
	for _, v := range poissonFloats(f) {
		largest = max(largest, math.Abs(v))
	}
	return largest
}

// lineTransformPlan carries the chirp-z pieces one transform length
// needs for the orthonormal DST-I (sine) or DCT-I (cosine) the Poisson
// solves apply along a line, plus the buffers the per-line work reuses.
//
// Both transforms are the real or imaginary part of a sum of complex
// exponentials over the half-frequency grid the trigonometric identity
// (j+1)(k+1) = ((j+1)² + (k+1)² − (j−k)²)/2 exposes: a per-sample
// phase, a per-bin rotation and an even chirp kernel, so the sum over
// the bins is a circular convolution of length the smallest power of
// two ≥ 2n−1. The padded full-length route each line used to take ran
// a convolution of roughly twice that length to reach the same bins.
type lineTransformPlan struct {
	n    int
	m    int
	sine bool
	// pre weights the samples, post rotates the bins, kern is the
	// kernel's forward transform on the convolution grid.
	pre  []complex128
	post []complex128
	kern []complex128
	norm float64
	// buf carries one line's spectrum through the convolution; the plan
	// is used from one goroutine and the solves are single-threaded, so
	// the buffer is fully rewritten on every apply.
	buf []complex128
}

// newLineTransformPlan builds the plan for n points. The chirp angles
// are reduced modulo the kernel's period in exact integer arithmetic
// before the transcendental call, so no angle grows with n.
func newLineTransformPlan(n int, sine bool) *lineTransformPlan {
	p := &lineTransformPlan{n: n, sine: sine}
	m := 1
	for m < 2*n-1 {
		m <<= 1
	}
	p.m = m
	p.pre = make([]complex128, n)
	p.post = make([]complex128, n)
	p.kern = make([]complex128, m)
	p.buf = make([]complex128, m)
	if sine {
		// DST-I: out[k] = norm·Im[ post_k · Σ_j x_j·pre_j·e^{−iθ(j−k)²/2} ]
		// with θ = π/(n+1), pre_j = e^{iθ(j + j²/2)},
		// post_k = e^{iθ(k+1 + k²/2)}.
		period := 4 * (n + 1)
		theta := math.Pi / float64(n+1)
		for j := range n {
			t := (2*j + j*j) % period
			p.pre[j] = polar(1, theta*float64(t)/2)
		}
		for k := range n {
			t := (2*(k+1) + k*k) % period
			p.post[k] = polar(1, theta*float64(t)/2)
		}
		for l := range n {
			t := (l * l) % period
			w := polar(1, -theta*float64(t)/2)
			p.kern[l] = w
			if l != 0 {
				p.kern[m-l] = w
			}
		}
		p.norm = math.Sqrt(2 / float64(n+1))
	} else {
		// DCT-I: out[k] = norm·Re[ w_k·e^{iθk²/2} ·
		// Σ_j x_j·w_j·e^{iθj²/2}·e^{−iθ(j−k)²/2} ] with θ = π/(n−1)
		// and w the half weight on the endpoints.
		period := 4 * (n - 1)
		theta := math.Pi / float64(n-1)
		for j := range n {
			t := (j * j) % period
			w := 1.0
			if j == 0 || j == n-1 {
				w = math.Sqrt2 / 2
			}
			p.pre[j] = complex(w, 0) * polar(1, theta*float64(t)/2)
		}
		for k := range n {
			t := (k * k) % period
			w := 1.0
			if k == 0 || k == n-1 {
				w = math.Sqrt2 / 2
			}
			p.post[k] = complex(w, 0) * polar(1, theta*float64(t)/2)
		}
		for l := range n {
			t := (l * l) % period
			w := polar(1, -theta*float64(t)/2)
			p.kern[l] = w
			if l != 0 {
				p.kern[m-l] = w
			}
		}
		p.norm = math.Sqrt(2 / float64(n-1))
	}
	// The kernel's spectrum depends only on the plan; it is the same
	// forward transform every line multiplies against.
	transform(p.kern, -1)
	return p
}

// polar returns the unit complex number of the given angle.
// math.Sincos returns exactly the pair math.Cos and math.Sin produce
// (verified bit-for-bit), so the value is unchanged.
func polar(magnitude, angle float64) complex128 {
	s, c := math.Sincos(angle)
	return complex(magnitude*c, magnitude*s)
}

// apply transforms src into dst; both have length p.n and must not
// alias. The convolution runs on the plan's reused spectrum buffer,
// which every call overwrites completely before it is read.
func (p *lineTransformPlan) apply(dst, src []float64) {
	n, m := p.n, p.m
	buf := p.buf
	if p.sine && n == 1 {
		dst[0] = src[0]
		return
	}
	for j := range n {
		buf[j] = complex(src[j], 0) * p.pre[j]
	}
	clear(buf[n:m])
	transform(buf, -1)
	for i := range m {
		buf[i] *= p.kern[i]
	}
	transform(buf, +1)
	scale := complex(1/float64(m), 0)
	if p.sine {
		for k := range n {
			v := buf[k] * scale * p.post[k]
			dst[k] = p.norm * imag(v)
		}
		return
	}
	for k := range n {
		v := buf[k] * scale * p.post[k]
		dst[k] = p.norm * real(v)
	}
}

// transformBlock applies the named self-inverse orthonormal transform
// (DST-I when sine, DCT-I when not) to every row and then every column
// of the rows×cols block of f whose top-left corner sits at (row0,
// col0), returning the flat spectrum. One plan per axis length serves
// every line of that axis.
func transformBlock(f *core.Array, row0, col0, rows, cols int, sine bool) ([]float64, error) {
	srcVals := poissonFloats(f)
	srcCols := f.Shape()[1]
	work := make([]float64, rows*cols)
	longest := max(rows, cols)
	line := make([]float64, longest)
	out := make([]float64, longest)
	// poissonFloats materialises a strided source in element order; a
	// raw payload read would keep the physical order instead.
	rowPlan := newLineTransformPlan(cols, sine)
	for r := range rows {
		for c := range cols {
			line[c] = srcVals[(r+row0)*srcCols+(c+col0)]
		}
		rowPlan.apply(out[:cols], line[:cols])
		copy(work[r*cols:(r+1)*cols], out[:cols])
	}
	colPlan := newLineTransformPlan(rows, sine)
	for c := range cols {
		for r := range rows {
			line[r] = work[r*cols+c]
		}
		colPlan.apply(out[:rows], line[:rows])
		for r := range rows {
			work[r*cols+c] = out[r]
		}
	}
	return work, nil
}

// inverseBlock applies the same self-inverse transform to a flat
// spectrum, the return trip of transformBlock.
func inverseBlock(spectrum []float64, rows, cols int, sine bool) ([]float64, error) {
	arr, err := core.FromFloats(spectrum, rows, cols)
	if err != nil {
		return nil, err
	}
	return transformBlock(arr, 0, 0, rows, cols, sine)
}