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

package integrate

import (
	"math"
	"sync"

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

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

// Two-dimensional evolution equations on the rectangle, the
// higher-rank siblings of IntegrateHeat1D and IntegrateWave1D. The
// grid is the rank-2 shape of the initial state: row r samples
// y = r·dy, column c samples x = c·dx, the boundary ring is held
// fixed, and the interior carries the dynamics.

// pde2dValidate checks the shared rectangle arguments and returns the
// grid shape.
func pde2dValidate(name string, u0 *core.Array, dx, dy, tFinal, dt float64, samples int) (rows, cols int, err error) {
	if u0.NDim() != 2 {
		return 0, 0, base.Errf("%s: the initial state must be rank 2, got shape %s", name, base.ShapeText(u0.Shape()))
	}
	if u0.Dtype() == core.Complex {
		return 0, 0, base.Errf("%s: complex initial states are not supported", name)
	}
	rows, cols = u0.Shape()[0], u0.Shape()[1]
	if rows < 3 || cols < 3 {
		return 0, 0, base.Errf("%s: the grid must be at least 3×3 to hold interior points, got %d×%d", name, rows, cols)
	}
	if !(dx > 0) || !(dy > 0) {
		return 0, 0, base.Errf("%s: the spacings must be positive, got %g and %g", name, dx, dy)
	}
	if !(tFinal > 0) || !(dt > 0) {
		return 0, 0, base.Errf("%s: tFinal and dt must be positive, got %g and %g", name, tFinal, dt)
	}
	// As in the 1-D validation: a dt this far below tFinal cannot be
	// honoured, and the wrapped step count would disable the bound
	// silently.
	if tFinal/dt > 1e12 {
		return 0, 0, base.Errf("%s: dt = %g asks for more than 1e12 steps over %g", name, dt, tFinal)
	}
	if samples < 2 {
		return 0, 0, base.Errf("%s: at least two samples are needed, got %d", name, samples)
	}
	// A non-finite entry flows through the stencils and the solvers'
	// zero-pivot guards compare false against NaN, so it is refused
	// up front.
	for i := range u0.Len() {
		if v := u0.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
			return 0, 0, base.Errf("%s: the initial state holds the non-finite value %g at %d", name, v, i)
		}
	}
	return rows, cols, nil
}

// fromSlice wraps a float slice of known length as an array.
func fromSlice(vals []float64, n int) *core.Array {
	out := core.New(core.Float, n)
	copy(out.RawFloats(), vals[:n])
	return out
}

// IntegrateHeat2D evolves u_t = κ·Δu on the rectangle by
// Peaceman-Rachford alternating direction implicit steps: one
// half-step implicit in x against an explicit y neighbour sum, one
// implicit in y against an explicit x neighbour sum, each row and
// column a tridiagonal solve through SolveTridiagonal. The scheme is
// second order in space and time and unconditionally stable, so no
// step-size refusal stands between the caller and a coarse first
// look. The boundary ring is held at the four constant edge values.
// The return is a (samples × rows × cols) array: the initial state,
// then the state after each stored interval, the final state forced
// into the last sample.
func IntegrateHeat2D(u0 *core.Array, kappa, dx, dy, tFinal, dt float64, samples int,
	boundBottom, boundTop, boundLeft, boundRight float64) (*core.Array, error) {
	const name = "IntegrateHeat2D"
	rows, cols, err := pde2dValidate(name, u0, dx, dy, tFinal, dt, samples)
	if err != nil {
		return nil, err
	}
	if !(kappa > 0) || math.IsInf(kappa, 0) {
		return nil, base.Errf("%s: the diffusivity must be positive, got %g", name, kappa)
	}
	// The four boundary constants enter the explicit neighbour sums
	// every step: a non-finite one would flow through the stencils and
	// the solves and publish an all-NaN history with no error.
	if math.IsNaN(boundBottom) || math.IsInf(boundBottom, 0) ||
		math.IsNaN(boundTop) || math.IsInf(boundTop, 0) ||
		math.IsNaN(boundLeft) || math.IsInf(boundLeft, 0) ||
		math.IsNaN(boundRight) || math.IsInf(boundRight, 0) {
		return nil, base.Errf("%s: the boundary values must be finite, got bottom %g, top %g, left %g, right %g",
			name, boundBottom, boundTop, boundLeft, boundRight)
	}
	steps, h := pdeSchedule(tFinal, dt, samples)
	rx := kappa * h / (2 * dx * dx) // the implicit half-step weight in x
	ry := kappa * h / (2 * dy * dy) // ...and in y

	u := make([]float64, rows*cols)
	copy(u, denseFloats(u0))
	every := steps / (samples - 1)
	out := core.New(core.Float, samples, rows, cols)
	into := out.RawFloats()
	// Sample 0 is the initial state exactly, before any ring
	// enforcement touches the working state.
	copy(into, u)
	written := 1
	// The working state's horizontal ring rows carry the boundary
	// constants from here on, so the stencils read their up and down
	// neighbours straight out of the payload: a ring slot holds exactly
	// the value the boundary test it replaces would have substituted.
	// The step loop below rewrites both rows after every half-step,
	// so the invariant holds for every step, the first included.
	for c := range cols {
		u[c] = boundBottom
		u[(rows-1)*cols+c] = boundTop
	}

	star := make([]float64, rows*cols)
	var stepErr error
	var errMu sync.Mutex
	// The implicit diagonals are constants of the scheme, one set per
	// orientation, built once with exactly the values the solves read.
	lowerX := make([]float64, cols-3)
	upperX := make([]float64, cols-3)
	diagX := make([]float64, cols-2)
	for i := range lowerX {
		lowerX[i] = -rx
		upperX[i] = -rx
	}
	for i := range diagX {
		diagX[i] = 1 + 2*rx
	}
	lowerY := make([]float64, rows-3)
	upperY := make([]float64, rows-3)
	diagY := make([]float64, rows-2)
	for i := range lowerY {
		lowerY[i] = -ry
		upperY[i] = -ry
	}
	for i := range diagY {
		diagY[i] = 1 + 2*ry
	}
	// Per-solve line scratch, one set per pdeStepFloor lines of each
	// orientation. engine.ParallelMin invokes a sweep closure once per
	// chunk; in parallel the chunk starts differ by at least
	// pdeStepFloor, so two concurrent chunks never share a set, and a
	// sweep the engine runs inline reuses set zero serially.
	triX := make([]triScratch, (rows-2-1)/pdeStepFloor+1)
	triY := make([]triScratch, (cols-2-1)/pdeStepFloor+1)
	for i := range triX {
		triSized(&triX[i], cols-2, cols)
	}
	for i := range triY {
		triSized(&triY[i], rows-2, rows)
	}
	for s := 1; s <= steps; s++ {
		// Half-step one: implicit in x along every interior row. Rows are
		// independent (a row's solve reads its own and its neighbours'
		// values and writes the star row), so the sweep is partitioned
		// over them; each chunk carries its own right-hand side and
		// solve scratch, and the arithmetic of a row is untouched.
		engine.ParallelMin(rows-2, pdeStepFloor, func(rs, re int) {
			tri := &triX[rs/pdeStepFloor]
			rhs := tri.rhs
			for r := rs + 1; r < re+1; r++ {
				// The row's own values and its two neighbours are slices
				// of the payload, so the stencil reads them by offset
				// instead of rebuilding the flat index per element. The
				// horizontal ring rows hold the boundary constants, so
				// the neighbours need no edge test.
				row := u[r*cols : (r+1)*cols]
				up := u[(r+1)*cols : (r+2)*cols]
				down := u[(r-1)*cols : r*cols]
				for c := range cols {
					rhs[c] = row[c] + ry*(up[c]-2*row[c]+down[c])
				}
				// Boundary neighbours of the implicit solve: the row's
				// own left and right ring values enter the right side.
				rhs[1] += rx * boundLeft
				rhs[cols-2] += rx * boundRight
				// The elimination writes the row's interior straight
				// into the star state; the ring values follow.
				dst := star[r*cols+1 : r*cols+cols-1]
				if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerX, diagX, upperX, rhs[1:cols-1]); serr != nil {
					errMu.Lock()
					if stepErr == nil {
						stepErr = base.Errf("%s: %w", name, serr)
					}
					errMu.Unlock()
					return
				}
				star[r*cols] = boundLeft
				star[r*cols+cols-1] = boundRight
			}
		})
		if stepErr != nil {
			return nil, stepErr
		}
		for c := range cols {
			star[c] = boundBottom
			star[(rows-1)*cols+c] = boundTop
		}
		// Half-step two: implicit in y down every interior column,
		// with the explicit x neighbours from the star state. Columns are
		// independent the same way the rows were, and the walk takes them
		// two at a time: the star state's vertical ring columns hold the
		// boundary constants, so both lanes read their neighbours
		// straight out of the payload, and a pair shares its loads, so
		// each line of the state is fetched once per pair instead of
		// three times across three column passes.
		engine.ParallelMin(cols-2, pdeStepFloor, func(cs, ce int) {
			tri := &triY[cs/pdeStepFloor]
			rhs := tri.rhs
			aux := tri.aux
			dst := tri.dst
			for c := cs + 1; c < ce+1; c += 2 {
				if c+1 < ce+1 {
					// Columns c and c+1: one pass over the rows builds
					// both right sides, the shared neighbour loads in
					// registers.
					for r := range rows {
						base := r * cols
						l := star[base+c-1]
						wm := star[base+c]
						e := star[base+c+1]
						rhs[r] = wm + rx*(e-2*wm+l)
						aux[r] = e + rx*(star[base+c+2]-2*e+wm)
					}
					rhs[1] += ry * boundBottom
					rhs[rows-2] += ry * boundTop
					if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerY, diagY, upperY, rhs[1:rows-1]); serr != nil {
						errMu.Lock()
						if stepErr == nil {
							stepErr = base.Errf("%s: %w", name, serr)
						}
						errMu.Unlock()
						return
					}
					u[c] = boundBottom
					u[(rows-1)*cols+c] = boundTop
					// The column's interior comes straight out of the
					// solve's payload: the values are the ones the
					// accessor read.
					for r := 1; r < rows-1; r++ {
						u[r*cols+c] = dst[r-1]
					}
					aux[1] += ry * boundBottom
					aux[rows-2] += ry * boundTop
					if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerY, diagY, upperY, aux[1:rows-1]); serr != nil {
						errMu.Lock()
						if stepErr == nil {
							stepErr = base.Errf("%s: %w", name, serr)
						}
						errMu.Unlock()
						return
					}
					u[c+1] = boundBottom
					u[(rows-1)*cols+c+1] = boundTop
					for r := 1; r < rows-1; r++ {
						u[r*cols+c+1] = dst[r-1]
					}
					continue
				}
				// The odd tail column, when the chunk ends on one.
				for r := range rows {
					base := r * cols
					wm := star[base+c]
					rhs[r] = wm + rx*(star[base+c+1]-2*wm+star[base+c-1])
				}
				rhs[1] += ry * boundBottom
				rhs[rows-2] += ry * boundTop
				if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerY, diagY, upperY, rhs[1:rows-1]); serr != nil {
					errMu.Lock()
					if stepErr == nil {
						stepErr = base.Errf("%s: %w", name, serr)
					}
					errMu.Unlock()
					return
				}
				u[c] = boundBottom
				u[(rows-1)*cols+c] = boundTop
				for r := 1; r < rows-1; r++ {
					u[r*cols+c] = dst[r-1]
				}
			}
		})
		if stepErr != nil {
			return nil, stepErr
		}
		for r := range rows {
			u[r*cols] = boundLeft
			u[r*cols+cols-1] = boundRight
		}
		if s%every == 0 && written < samples {
			copy(into[written*rows*cols:(written+1)*rows*cols], u)
			written++
		}
	}
	copy(into[(samples-1)*rows*cols:], u)
	return out, nil
}

// IntegrateWave2D evolves u_tt = c²·Δu on the rectangle with the
// boundary ring held at zero, by the explicit central-difference
// stencil from the velocity Verlet family the 1-D wave solver uses:
// second order in space and time, with the CFL budget
// c·dt·sqrt(1/dx² + 1/dy²) ≤ 1 enforced as an error, because the
// explicit stencil has no honest answer past it. The return contract
// mirrors IntegrateHeat2D.
func IntegrateWave2D(u0, v0 *core.Array, c, dx, dy, tFinal, dt float64, samples int) (*core.Array, error) {
	const name = "IntegrateWave2D"
	rows, cols, err := pde2dValidate(name, u0, dx, dy, tFinal, dt, samples)
	if err != nil {
		return nil, err
	}
	if v0.NDim() != 2 || v0.Shape()[0] != rows || v0.Shape()[1] != cols {
		return nil, base.Errf("%s: the velocity must be rank 2 on the same grid, got shape %s", name, base.ShapeText(v0.Shape()))
	}
	if v0.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex velocities are not supported", name)
	}
	// As for u0: a non-finite velocity enters the Taylor start and
	// poisons the three-level stencil without an error.
	for i := range v0.Len() {
		if v := v0.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, base.Errf("%s: the velocity holds the non-finite value %g at %d", name, v, i)
		}
	}
	if !(c > 0) {
		return nil, base.Errf("%s: the wave speed must be positive, got %g", name, c)
	}
	cfl := c * dt * math.Sqrt(1/(dx*dx)+1/(dy*dy))
	if cfl > 1 {
		return nil, base.Errf("%s: the CFL number %g exceeds 1 (c = %g, dt = %g, dx = %g, dy = %g); the explicit stencil is unstable there",
			name, cfl, c, dt, dx, dy)
	}
	steps, h := pdeSchedule(tFinal, dt, samples)
	lapW := (c * h) * (c * h)
	// The stencil's spacings are constants of the grid; squaring them
	// once per step is the same product the element loop evaluated.
	dx2 := dx * dx
	dy2 := dy * dy

	prev := make([]float64, rows*cols) // u^{s-1}
	cur := make([]float64, rows*cols)  // u^{s}
	next := make([]float64, rows*cols) // u^{s+1}, the buffer being written
	copy(cur, denseFloats(u0))
	every := steps / (samples - 1)
	out := core.New(core.Float, samples, rows, cols)
	into := out.RawFloats()
	// Slot 0 is the initial state at t = 0; slot j is the state after
	// j·every steps, at the published time j·tFinal/(samples-1), exactly
	// like IntegrateHeat2D.
	copy(into, cur)
	// From there on the boundary ring is held at zero, so clear it in the
	// working buffer the initial state came in (the other two are born
	// zeroed): otherwise a recycled buffer would carry the caller's ring
	// back into the stencil every third step.
	for r := range rows {
		cur[r*cols], cur[r*cols+cols-1] = 0, 0
	}
	for c := range cols {
		cur[c], cur[(rows-1)*cols+c] = 0, 0
	}
	written := 1
	// The three buffers cycle through (prev, cur, next), so every stencil
	// reads u^{s-1} and u^s while writing u^{s+1} into a third buffer: no
	// point is ever overwritten before its neighbours have read it. The
	// first iteration is the Taylor start
	// u¹ = u⁰ + h·v⁰ + h²c²/2·Δu⁰, second order, so the three-level
	// stencil starts honest, and it reads the untouched u⁰ (the initial
	// velocity and the Laplacian of u⁰), not a half-updated state.
	for s := 1; s <= steps; s++ {
		first := s == 1
		for r := 1; r < rows-1; r++ {
			// The row's neighbours are the slices above and below it, so
			// the stencil reads them by offset instead of rebuilding the
			// flat index per element. The terms keep their order.
			mid := cur[r*cols : (r+1)*cols]
			up := cur[(r+1)*cols : (r+2)*cols]
			down := cur[(r-1)*cols : r*cols]
			if first {
				// The Taylor start reads the untouched u⁰ and its
				// initial velocity, not a half-updated state.
				for cc := 1; cc < cols-1; cc++ {
					i := r*cols + cc
					lap := (down[cc]-2*mid[cc]+up[cc])/dy2 +
						(mid[cc-1]-2*mid[cc]+mid[cc+1])/dx2
					next[i] = mid[cc] + h*v0.FloatAt(i) + 0.5*lapW*lap
				}
				continue
			}
			for cc := 1; cc < cols-1; cc++ {
				i := r*cols + cc
				lap := (down[cc]-2*mid[cc]+up[cc])/dy2 +
					(mid[cc-1]-2*mid[cc]+mid[cc+1])/dx2
				next[i] = 2*mid[cc] - prev[i] + lapW*lap
			}
		}
		prev, cur, next = cur, next, prev
		if s%every == 0 && written < samples {
			copy(into[written*rows*cols:(written+1)*rows*cols], cur)
			written++
		}
	}
	copy(into[(samples-1)*rows*cols:], cur)
	return out, nil
}

// pdeStepFloor is the per-worker floor for a 2-D step sweep. One item is
// one tridiagonal elimination over a stencil row, measured at roughly
// half a microsecond on a 32-line grid, so a worker needs about eight
// of them before the split pays for its spawn. The floor is a constant
// for that reason, not a length-scaled budget. The line scratch is
// indexed by start/pdeStepFloor on the same constant: parallel chunk
// starts differ by at least the floor, so concurrent chunks never share
// a scratch set.
const pdeStepFloor = 8