// 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"
)

// Advection in one space dimension, the transport siblings of
// IntegrateHeat1D and IntegrateWave1D: u_t + a·u_x = 0 and the
// advection-diffusion equation u_t + a·u_x = D·u_xx. Transport is
// where the cheap stencils fail in public: the first-order upwind
// flux is monotone but smears a front every step it takes, and the
// centred flux that heat uses is oscillatory or worse here. The
// middle ground is a flux-limited scheme: an upwind flux whose face
// value is raised toward the third-order one by a slope limiter that
// switches itself off across discontinuities, keeping the scheme
// total variation diminishing.
//
// The limiter is Koren's third-order one, φ(θ) = max(0, min(2θ,
// (2+θ)/3, 2)), stated here as the choice. The face value is the
// Sweby flux-limited form u_upwind + ½φ(θ)·(1−|ν|)·Δ, with ν =
// |a|·dt/dx the CFL number and Δ the forward difference across the
// face: the (1−|ν|) factor is what makes the explicit update total
// variation diminishing in the Sweby sense for CFL ≤ 1, and the
// third-order branch of φ is what keeps smooth profiles sharp (a
// marginally small overshoot past the strict Sweby bound survives as
// roundoff-scale noise). The grid convention is the house one: u0
// carries the interior cells with dx = L/(n+1), and boundL, boundR
// are the fixed values of the ghost cells on the ends. The face
// adjacent to the inflow end takes its prescribed value directly
// (first order there); on every other face, including the outflow
// end, the limiter runs at full strength with the prescribed ghost
// value entering its smoothness ratio, and a window too short for a
// three-cell stencil forces first order.

// korenSlope returns the Koren-limited normalised slope φ(θ) for the
// smoothness ratio θ: the three branches are the Sweby bounds with
// the third-order diagonal (2+θ)/3.
func korenSlope(theta float64) float64 {
	phi := min(2*theta, (2+theta)/3, 2)
	if phi < 0 {
		return 0
	}
	return phi
}

// advectFaceFlux returns the numerical flux a·u_face across the face
// between cells ul (left) and ur (right), with the second neighbours
// ull (left of ul) and urr (right of ur) feeding the limiter and nu
// the CFL number feeding the time factor. The wind decides the side
// the face value is reconstructed from; with limited false the flux
// is plain first-order upwind.
func advectFaceFlux(a, ull, ul, ur, urr, nu float64, limited bool) float64 {
	d := ur - ul
	if a >= 0 {
		if !limited || d == 0 {
			return a * ul
		}
		theta := (ul - ull) / d
		return a * (ul + 0.5*korenSlope(theta)*(1-nu)*d)
	}
	if !limited || d == 0 {
		return a * ur
	}
	theta := (ur - urr) / d
	return a * (ur - 0.5*korenSlope(theta)*(1-nu)*d)
}

// advectStep advances dst from src by one explicit step of width h
// of the conservative update u −= (h/dx)·(F₊ − F₋), the faces built
// from src with the ghost values boundL and boundR on the ends. dst
// and src must not alias. Face j sits between cell j−1 and cell j,
// so face 0 borders the left ghost and face n the right one.
func advectStep(dst, src []float64, a, dx, h, boundL, boundR float64, limited bool, faces []float64) {
	n := len(src)
	lambda := h / dx
	nu := math.Abs(a) * lambda
	for j := range n + 1 {
		switch {
		case j == 0:
			// The inflow face takes the prescribed ghost value; on an
			// outflow left end the upwind cell is cell 0, whose
			// limited face value reaches one cell into the interior.
			if a >= 0 {
				faces[j] = a * boundL
			} else if n < 2 || !limited {
				faces[j] = a * src[0]
			} else {
				faces[j] = advectFaceFlux(a, boundL, boundL, src[0], src[1], nu, true)
			}
		case j == n:
			if a >= 0 {
				if n < 2 || !limited {
					faces[j] = a * src[n-1]
				} else {
					faces[j] = advectFaceFlux(a, src[n-2], src[n-1], boundR, boundR, nu, true)
				}
			} else {
				faces[j] = a * boundR
			}
		default:
			ull := boundL
			if j >= 2 {
				ull = src[j-2]
			}
			urr := boundR
			if j <= n-2 {
				urr = src[j+1]
			}
			faces[j] = advectFaceFlux(a, ull, src[j-1], src[j], urr, nu, limited)
		}
	}
	for i := range n {
		dst[i] = src[i] - lambda*(faces[i+1]-faces[i])
	}
}

// advectValidate checks the shared input contract of the transport
// solvers and returns the cell count. The grid, step bound, sample
// and finiteness gates are pdeValidate's; transport adds a finite
// speed and finite ghost values, and enforces the CFL budget
// |a|·dt/dx ≤ 1 the explicit update needs, exactly like the wave
// solver enforces its own.
func advectValidate(name string, u0 *core.Array, a, dx, tFinal, dt float64, samples int, boundL, boundR float64) (int, error) {
	n, err := pdeValidate(name, u0, dx, tFinal, dt, samples)
	if err != nil {
		return 0, err
	}
	if math.IsNaN(a) || math.IsInf(a, 0) {
		return 0, base.Errf("%s: the transport speed must be finite, got %g", name, a)
	}
	if math.IsNaN(boundL) || math.IsInf(boundL, 0) || math.IsNaN(boundR) || math.IsInf(boundR, 0) {
		return 0, base.Errf("%s: the ghost values must be finite, got %g and %g", name, boundL, boundR)
	}
	cfl := math.Abs(a * dt / dx)
	if cfl > 1 {
		return 0, base.Errf("%s: CFL violated, |a·dt/dx| = %.3g > 1", name, cfl)
	}
	return n, nil
}

// IntegrateAdvection1D evolves u_t + a·u_x = 0 over the grid of u0
// from t = 0 to tFinal in equal steps of at most dt, and returns the
// (samples, n) array of interior states evenly spaced in time,
// endpoints included, exactly like IntegrateHeat1D. The flux is the
// Koren-limited upwind one described at the top of the file: total
// variation diminishing under CFL ≤ 1, third-order at smooth faces
// and first order next to the inflow boundary, so a front is carried
// sharply where plain upwind would smear it away.
func IntegrateAdvection1D(u0 *core.Array, a, dx, tFinal, dt float64, samples int, boundL, boundR float64) (*core.Array, error) {
	const name = "IntegrateAdvection1D"
	n, err := advectValidate(name, u0, a, dx, tFinal, dt, samples, boundL, boundR)
	if err != nil {
		return nil, err
	}
	return advectRun(name, u0, a, dx, tFinal, dt, samples, boundL, boundR, n, true)
}

// IntegrateUpwindAdvection1D evolves the same equation with the
// plain first-order upwind flux: monotone under CFL ≤ 1 (no new
// extrema, ever) and diffuse, the baseline the limited scheme is
// measured against. The return contract mirrors IntegrateAdvection1D.
func IntegrateUpwindAdvection1D(u0 *core.Array, a, dx, tFinal, dt float64, samples int, boundL, boundR float64) (*core.Array, error) {
	const name = "IntegrateUpwindAdvection1D"
	n, err := advectValidate(name, u0, a, dx, tFinal, dt, samples, boundL, boundR)
	if err != nil {
		return nil, err
	}
	return advectRun(name, u0, a, dx, tFinal, dt, samples, boundL, boundR, n, false)
}

// advectRun is the shared stepping loop of the two pure-transport
// solvers: fixed steps on the pdeSchedule grid, states sampled every
// steps/(samples−1) steps with the final state forced into the last
// sample.
func advectRun(name string, u0 *core.Array, a, dx, tFinal, dt float64, samples int, boundL, boundR float64, n int, limited bool) (*core.Array, error) {
	u := make([]float64, n)
	for i := range n {
		u[i] = u0.FloatAt(i)
	}
	steps, h := pdeSchedule(tFinal, dt, samples)
	every := steps / (samples - 1)
	out := core.New(core.Float, samples, n)
	into := out.RawFloats()
	copy(into, u)
	written := 1
	faces := make([]float64, n+1)
	scratch := make([]float64, n)
	for s := 1; s <= steps; s++ {
		advectStep(scratch, u, a, dx, h, boundL, boundR, limited, faces)
		u, scratch = scratch, u
		if s%every == 0 && written < samples {
			copy(into[written*n:(written+1)*n], u)
			written++
		}
	}
	copy(into[(samples-1)*n:], u)
	return out, nil
}

// IntegrateAdvectionDiffusion1D evolves u_t + a·u_x = D·u_xx over the
// grid of u0 with the Dirichlet ghost values boundL and boundR. Each
// step combines the Koren-limited advection flux, advanced
// explicitly, with the Crank-Nicolson second-difference diffusion the
// heat solver runs through the shared tridiagonal solve, so the
// composition is first order in time and second order in space and
// the diffusion side is unconditionally stable. The explicit
// advection still answers for its own CFL budget |a|·dt/dx ≤ 1 and
// the step is refused past it. With a = 0 the scheme reduces exactly
// to IntegrateHeat1D.
func IntegrateAdvectionDiffusion1D(u0 *core.Array, a, kappa, dx, tFinal, dt float64, samples int, boundL, boundR float64) (*core.Array, error) {
	const name = "IntegrateAdvectionDiffusion1D"
	n, err := advectValidate(name, u0, a, dx, tFinal, dt, samples, boundL, boundR)
	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)
	}
	u := make([]float64, n)
	for i := range n {
		u[i] = u0.FloatAt(i)
	}
	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 implicit left side I − r/2·A is a constant of the scheme,
	// built once exactly as IntegrateHeat1D builds it, strictly
	// diagonally dominant for every positive r like the heat system.
	for i := range n {
		diag[i] = 1 + r
		if i < n-1 {
			lower[i] = -r / 2
			upper[i] = -r / 2
		}
	}
	// The trapezoidal weight and the elimination scratch are constants
	// of one solve: every step refills the same right side and the
	// solution is written straight into the working state.
	half := r / 2
	var tri triScratch
	triSized(&tri, n, n)

	every := steps / (samples - 1)
	out := core.New(core.Float, samples, n)
	into := out.RawFloats()
	copy(into, u)
	written := 1
	faces := make([]float64, n+1)
	advected := make([]float64, n)
	rhs := tri.rhs
	for s := 1; s <= steps; s++ {
		// Explicit advection sub-step on the limited fluxes.
		advectStep(advected, u, a, dx, h, boundL, boundR, true, faces)
		// Crank-Nicolson diffusion sub-step on the advected state,
		// the Dirichlet neighbours entering as known data on both
		// sides, in the heat solver's own arithmetic.
		for i := range n {
			um, up := boundL, boundR
			if i > 0 {
				um = advected[i-1]
			}
			if i < n-1 {
				up = advected[i+1]
			}
			rhs[i] = advected[i] + half*(um-2*advected[i]+up)
			if i == 0 {
				rhs[i] += half * boundL
			}
			if i == n-1 {
				rhs[i] += half * boundR
			}
		}
		if serr := base.TriSolve(u, tri.cp, tri.dp, lower, diag, upper, rhs); serr != nil {
			return nil, base.Errf("%s: %w", name, serr)
		}
		if s%every == 0 && written < samples {
			copy(into[written*n:(written+1)*n], u)
			written++
		}
	}
	copy(into[(samples-1)*n:], u)
	return out, nil
}