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

// Symplectic integration for separable Hamiltonian systems, where the
// energy splits as H(q, p) = T(p) + V(q): Newtonian mechanics, N-body
// gravity, molecular dynamics. The adaptive Runge-Kutta pair that
// drives IntegrateODE is accurate per step but dissipates energy
// systematically, so a two-hundred-period orbit spirals inward or
// outward; the leapfrog structure below is symplectic, which means it
// conserves a shadow Hamiltonian exactly and keeps the true energy
// oscillating in a bounded band forever. That long-time fidelity, not
// per-step accuracy, is what separates integrators for celestial
// mechanics and molecular dynamics from general-purpose ones.
//
// The scheme is velocity Verlet, a kick-drift-kick leapfrog of second
// order: half a momentum kick, a full position drift, another half
// kick with the force at the new position. Unit masses are assumed
// (p is the velocity); scale the momentum by the masses beforehand or
// fold them into the acceleration.

// IntegrateVerlet integrates a separable Hamiltonian system with unit
// masses over an even time grid: accel returns the acceleration
// −∂V/∂q at a position, q0 and p0 are the initial position and
// momentum (velocity), and steps fixes the number of equal steps, so
// the i-th returned pair sits at t0 + i·h with h = (t1−t0)/steps.
// positions[0] is q0 and momenta[0] is p0. The state may be float64
// or float32, read through per-element accessors, and the returned
// arrays are float64. The step size stays fixed by design:
// adaptivity would destroy the symplectic property the method exists
// for. A negative or zero-length acceleration vector, a mismatched
// pair, a complex or int state, or an empty state is an error.
func IntegrateVerlet(accel func(q *core.Array) (*core.Array, error),
	t0, t1 float64, q0, p0 *core.Array, steps int) (positions, momenta []*core.Array, err error) {
	if steps <= 0 {
		return nil, nil, base.Errf("IntegrateVerlet: steps must be ≥ 1, got %d", steps)
	}
	if q0.Dtype() == core.Complex || p0.Dtype() == core.Complex {
		return nil, nil, base.Errf("IntegrateVerlet: complex states are not supported")
	}
	// The whole integer class follows Int into the standing refusal:
	// bool and the narrow widths carry a discrete state, which has no
	// place in a continuous integrator, and the wording is Int's own.
	if integerState(q0.Dtype()) || integerState(p0.Dtype()) {
		return nil, nil, base.Errf("IntegrateVerlet: int states cannot integrate, use float or float32 states")
	}
	if q0.NDim() != 1 || p0.NDim() != 1 || q0.Len() != p0.Len() {
		return nil, nil, base.Errf("IntegrateVerlet: position and momentum must be vectors of equal length, got %s and %s",
			base.ShapeText(q0.Shape()), base.ShapeText(p0.Shape()))
	}
	n := q0.Len()
	if n == 0 {
		return nil, nil, base.Errf("IntegrateVerlet: the state must not be empty")
	}
	// Read the initial state element-wise: RawFloats backs float64
	// payloads only, so a float32 state would come through as nil.
	// A non-finite entry is refused up front: it would propagate
	// through every kick and drift silently.
	q := make([]float64, n)
	p := make([]float64, n)
	for i := range n {
		q[i] = q0.FloatAt(i)
		p[i] = p0.FloatAt(i)
		if math.IsNaN(q[i]) || math.IsInf(q[i], 0) {
			return nil, nil, base.Errf("IntegrateVerlet: q0 holds the non-finite value %g at %d", q[i], i)
		}
		if math.IsNaN(p[i]) || math.IsInf(p[i], 0) {
			return nil, nil, base.Errf("IntegrateVerlet: p0 holds the non-finite value %g at %d", p[i], i)
		}
	}
	a := make([]float64, n)
	// One cached read-only view serves every acceleration call: the
	// position slice is the run's own buffer, stable for the whole
	// integration, so the wrapper is built once per run.
	var views odeViews
	eval := func(x []float64, out []float64) error {
		v, err := accel(views.of(x))
		if err != nil {
			return base.Errf("IntegrateVerlet: %w", err)
		}
		if v.NDim() != 1 || v.Len() != n {
			return base.Errf("IntegrateVerlet: accel returned shape %s, want a vector of length %d",
				base.ShapeText(v.Shape()), n)
		}
		readVector(out, v)
		// Like RK4: a non-finite acceleration would flow through the
		// kicks silently, and the published trajectory would be NaN
		// with a nil error.
		for i := range n {
			if math.IsNaN(out[i]) || math.IsInf(out[i], 0) {
				return base.Errf("IntegrateVerlet: accel returned the non-finite value %g at coordinate %d", out[i], i)
			}
		}
		return nil
	}
	if err := eval(q, a); err != nil {
		return nil, nil, err
	}
	h := (t1 - t0) / float64(steps)
	positions = make([]*core.Array, steps+1)
	momenta = make([]*core.Array, steps+1)
	positions[0] = arrayFromVector(q)
	momenta[0] = arrayFromVector(p)
	for s := 1; s <= steps; s++ {
		// Kick, drift, kick: two half kicks bracket the drift, so the
		// force is evaluated once per step.
		for i := range n {
			p[i] += h / 2 * a[i]
			q[i] += h * p[i]
		}
		if err := eval(q, a); err != nil {
			return nil, nil, err
		}
		for i := range n {
			p[i] += h / 2 * a[i]
		}
		positions[s] = arrayFromVector(q)
		momenta[s] = arrayFromVector(p)
	}
	return positions, momenta, nil
}