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

package integrate

import (
	"math"

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

// Turnkey PDE evolution in one space dimension: the two
// equations half of physics reduces to, wrapped on machinery the
// library already owns. The heat equation runs Crank-Nicolson (the
// unconditionally stable trapezoidal rule) through the shared
// tridiagonal solver; the wave equation runs velocity Verlet on the
// second-order form, kick-drift-kick like the Hamiltonian integrator
// it is. Both return the trajectory sampled on a time grid,
// IntegrateODEPath-style.

// pdeValidate checks the shared input contract and returns the grid
// size.
func pdeValidate(name string, u0 *core.Array, dx, tFinal, dt float64, samples int) (int, error) {
	if u0.NDim() != 1 || u0.Len() == 0 {
		return 0, base.Errf("%s: the initial condition must be a non-empty rank-1 array, got shape %s",
			name, base.ShapeText(u0.Shape()))
	}
	if u0.Dtype() == core.Complex {
		return 0, base.Errf("%s: complex states are not supported", name)
	}
	if !(dx > 0) {
		return 0, base.Errf("%s: the grid spacing must be positive, got %g", name, dx)
	}
	if !(tFinal > 0) {
		return 0, base.Errf("%s: the integration time must be positive, got %g", name, tFinal)
	}
	if !(dt > 0) {
		return 0, base.Errf("%s: the time step must be positive, got %g", name, dt)
	}
	// A dt far below tFinal/1e12 cannot be honoured: the step count
	// would leave the int range on some platforms and wrap on others,
	// and the silently larger step would run past the wave equation's
	// CFL check, which runs on the requested dt.
	if tFinal/dt > 1e12 {
		return 0, base.Errf("%s: dt = %g asks for more than 1e12 steps over %g", name, dt, tFinal)
	}
	if samples < 2 {
		return 0, base.Errf("%s: at least two samples are needed, got %d", name, samples)
	}
	// A non-finite entry would flow through the stencil and the
	// tridiagonal solve's zero-pivot checks compare false against NaN,
	// publishing an all-NaN history with no error.
	for i := range u0.Len() {
		if v := u0.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
			return 0, base.Errf("%s: the initial condition holds the non-finite value %g at %d", name, v, i)
		}
	}
	return u0.Len(), nil
}

// pdeSchedule picks the step count and the actual step size for a
// requested dt. The count is rounded up to a multiple of the sampling
// interval, so every published time j·tFinal/(samples−1) is a step
// boundary and the last step lands on tFinal exactly: the returned
// samples are the evenly spaced interior states the documentation
// promises, not the states at multiples of dt.
func pdeSchedule(tFinal, dt float64, samples int) (steps int, h float64) {
	steps = max(int(math.Ceil(tFinal/dt)), samples-1)
	if rem := steps % (samples - 1); rem != 0 {
		steps += samples - 1 - rem
	}
	return steps, tFinal / float64(steps)
}

// IntegrateHeat1D evolves u_t = κ·u_xx over [0, L] discretised by the
// interior grid of u0 (n = u0.Len(), dx = L/(n+1)), from t = 0 to
// tFinal in equal steps of at most dt, holding the boundary values
// boundL and boundR (Dirichlet). It returns the (samples, n) array of
// interior states evenly spaced in time, endpoints included. Crank-
// Nicolson is stable for any dt; accuracy wants dt of a few dx²/κ.
func IntegrateHeat1D(u0 *core.Array, kappa, dx, tFinal, dt float64, samples int, boundL, boundR float64) (*core.Array, error) {
	const name = "IntegrateHeat1D"
	n, err := pdeValidate(name, u0, dx, 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 boundary values enter the right side every step: a non-finite
	// one would flow through the stencil and the solve and publish an
	// all-NaN history with no error.
	if math.IsNaN(boundL) || math.IsInf(boundL, 0) || math.IsNaN(boundR) || math.IsInf(boundR, 0) {
		return nil, base.Errf("%s: the boundary values must be finite, got %g and %g", name, boundL, boundR)
	}
	u := make([]float64, n)
	copy(u, denseFloats(u0))
	// Crank-Nicolson: (I − r/2·A)uⁿ⁺¹ = (I + r/2·A)uⁿ with A the
	// second-difference stencil and r = κ·h/dx² for the step h the
	// schedule actually takes; the Dirichlet neighbours enter the right
	// side through the stencil ends.
	steps, h := pdeSchedule(tFinal, dt, samples)
	r := kappa * h / (dx * dx)
	lower := make([]float64, n-1)
	diag := make([]float64, n)
	upper := make([]float64, n-1)
	// The left side carries I − r/2·A: the diagonal gains r (A's −2
	// times −r/2) and the off-diagonals stay −r/2. The system is
	// strictly diagonally dominant for every positive r, so the
	// elimination's pivots stay finite and non-zero.
	for i := range n {
		diag[i] = 1 + r
		if i < n-1 {
			lower[i] = -r / 2
			upper[i] = -r / 2
		}
	}

	every := steps / (samples - 1)
	out := make([]float64, samples*n)
	copy(out, u)
	written := 1
	// The right side and the elimination scratch are constants of one
	// solve: every step refills the same buffers, the kernel reads its
	// inputs without touching them, and the solution is written
	// straight into the working state the samples copy from.
	var tri triScratch
	triSized(&tri, n, n)
	rhs := tri.rhs
	// The trapezoidal weight is a constant of the scheme: one division
	// by two, the same value the expression inside the loop carried.
	half := r / 2
	for s := 1; s <= steps; s++ {
		for i := range n {
			um, up := boundL, boundR
			if i > 0 {
				um = u[i-1]
			}
			if i < n-1 {
				up = u[i+1]
			}
			rhs[i] = u[i] + half*(um-2*u[i]+up)
		}
		// The implicit side's boundary neighbours move across as known
		// data: the first and last rows only, in that order.
		rhs[0] += half * boundL
		rhs[n-1] += half * boundR
		if err := base.TriSolve(u, tri.cp, tri.dp, lower, diag, upper, rhs); err != nil {
			return nil, base.Errf("%s: %w", name, err)
		}
		if s%every == 0 && written < samples {
			copy(out[written*n:(written+1)*n], u)
			written++
		}
	}
	// The final state is the last sample whatever the grid remainder.
	copy(out[(samples-1)*n:], u)
	return core.FromFloats(out, samples, n)
}

// IntegrateWave1D evolves u_tt = c²·u_xx over [0, L] with the grid of
// u0 (dx = L/(n+1), Dirichlet ends held at zero) and the initial
// velocity v0, by velocity Verlet with fixed step dt. The CFL budget
// |c·dt/dx| ≤ 1 is a genuine stability requirement and is enforced as
// an error. The return contract mirrors IntegrateHeat1D.
func IntegrateWave1D(u0, v0 *core.Array, c, dx, tFinal, dt float64, samples int) (*core.Array, error) {
	const name = "IntegrateWave1D"
	n, err := pdeValidate(name, u0, dx, tFinal, dt, samples)
	if err != nil {
		return nil, err
	}
	if v0.NDim() != 1 || v0.Len() != n {
		return nil, base.Errf("%s: the initial velocity must match the state shape, got %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 flows through the Verlet kick
	// and poisons the trajectory without an error.
	for i := range n {
		if v := v0.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, base.Errf("%s: the initial velocity holds the non-finite value %g at %d", name, v, i)
		}
	}
	// The CFL ratio compares false against 1 when it is NaN, so a
	// non-finite speed is refused before the budget test.
	if math.IsNaN(c) || math.IsInf(c, 0) {
		return nil, base.Errf("%s: the wave speed must be finite, got %g", name, c)
	}
	cfl := math.Abs(c * dt / dx)
	if cfl > 1 {
		return nil, base.Errf("%s: CFL violated, |c·dt/dx| = %.3g > 1", name, cfl)
	}
	u := make([]float64, n)
	v := make([]float64, n)
	copy(u, denseFloats(u0))
	copy(v, denseFloats(v0))
	// The stencil's constants: the products are the ones the element
	// loop evaluated, built once per run.
	cc := c * c
	dx2 := dx * dx
	accel := func(dst, us []float64) {
		// The two ends take the fixed zero neighbour the boundaries
		// impose; the interior runs the same stencil over the real
		// neighbours, so the two tests leave the element loop. Every
		// term keeps the order the uniform loop evaluated.
		end := func(i int) {
			um, up := 0.0, 0.0
			if i > 0 {
				um = us[i-1]
			}
			if i < n-1 {
				up = us[i+1]
			}
			dst[i] = cc * (um - 2*us[i] + up) / dx2
		}
		end(0)
		for i := 1; i < n-1; i++ {
			dst[i] = cc * (us[i-1] - 2*us[i] + us[i+1]) / dx2
		}
		if n > 1 {
			end(n - 1)
		}
	}
	// Kick-drift-kick: the modified energy stays within (c·dt/dx)²/8
	// of the true one, which is why the wave equation keeps its shape.
	// The schedule's step never exceeds dt, so the CFL check above
	// bounds this one too.
	steps, h := pdeSchedule(tFinal, dt, samples)
	every := steps / (samples - 1)
	out := make([]float64, samples*n)
	copy(out, u)
	written := 1
	a := make([]float64, n)
	for s := 1; s <= steps; s++ {
		accel(a, u)
		for i := range n {
			v[i] += 0.5 * h * a[i]
		}
		for i := range n {
			u[i] += h * v[i]
		}
		accel(a, u)
		for i := range n {
			v[i] += 0.5 * h * a[i]
		}
		if s%every == 0 && written < samples {
			copy(out[written*n:(written+1)*n], u)
			written++
		}
	}
	copy(out[(samples-1)*n:], u)
	return core.FromFloats(out, samples, n)
}