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

package integrate

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

import (
	"math"
	"testing"
)

// TestIntegrateVerletHarmonicEnergy runs the harmonic oscillator over
// fifty periods: the true energy must stay inside a bounded band the
// whole time (the symplectic property) and the orbit must close on
// itself to second-order accuracy.
func TestIntegrateVerletHarmonicEnergy(t *testing.T) {
	const omega = 1.0
	accel := func(q *core.Array) (*core.Array, error) {
		return core.MulF(q, -omega*omega), nil
	}
	const periods = 50
	// With h per period the energy band sits at O((ωh)²/8)·E; 300
	// steps per period put it under the 1e-4 bar.
	stepsPerPeriod := 400
	q0 := mustFloats(t, []float64{1})
	p0 := mustFloats(t, []float64{0})
	positions, momenta, err := IntegrateVerlet(accel, 0, periods*2*math.Pi, q0, p0, periods*stepsPerPeriod)
	if err != nil {
		t.Fatalf("IntegrateVerlet: %v", err)
	}
	energy := func(q, p float64) float64 {
		return 0.5 * (p*p + omega*omega*q*q)
	}
	e0 := energy(positions[0].FloatAt(0), momenta[0].FloatAt(0))
	worst := 0.0
	for s := range positions {
		e := energy(positions[s].FloatAt(0), momenta[s].FloatAt(0))
		if d := math.Abs(e - e0); d > worst {
			worst = d
		}
	}
	// The band must stay narrow relative to the energy itself, over
	// fifty periods, where a non-symplectic scheme drifts away.
	if worst > 1e-4*e0 {
		t.Fatalf("energy drifted by %g (energy %g) over %d periods", worst, e0, periods)
	}
	last := positions[len(positions)-1].FloatAt(0)
	if math.Abs(last-1) > 1e-3 {
		t.Fatalf("after a whole number of periods q = %.10g, want ≈ 1", last)
	}
	// Second order: halving the step must cut the closing error by
	// roughly four.
	q1, _, err := IntegrateVerlet(accel, 0, 2*math.Pi, q0, p0, 20)
	if err != nil {
		t.Fatalf("IntegrateVerlet coarse: %v", err)
	}
	q2, _, err := IntegrateVerlet(accel, 0, 2*math.Pi, q0, p0, 40)
	if err != nil {
		t.Fatalf("IntegrateVerlet fine: %v", err)
	}
	eCoarse := math.Abs(q1[len(q1)-1].FloatAt(0) - 1)
	eFine := math.Abs(q2[len(q2)-1].FloatAt(0) - 1)
	if eFine > 0.35*eCoarse {
		t.Fatalf("convergence order looks wrong: coarse %g, fine %g", eCoarse, eFine)
	}
}

// TestIntegrateVerletPendulum checks a nonlinear system: the pendulum
// with E = p²/2 − cos q must keep its energy bounded, including near
// the separatrix where the force is far from linear.
func TestIntegrateVerletPendulum(t *testing.T) {
	accel := func(q *core.Array) (*core.Array, error) {
		return mustFloats(t, []float64{-math.Sin(q.FloatAt(0))}), nil
	}
	q0 := mustFloats(t, []float64{2.5})
	p0 := mustFloats(t, []float64{0})
	positions, momenta, err := IntegrateVerlet(accel, 0, 200, q0, p0, 20000)
	if err != nil {
		t.Fatalf("IntegrateVerlet: %v", err)
	}
	energyOf := func(q, p float64) float64 {
		return 0.5*p*p - math.Cos(q)
	}
	e0 := energyOf(positions[0].FloatAt(0), momenta[0].FloatAt(0))
	worst := 0.0
	for s := range positions {
		e := energyOf(positions[s].FloatAt(0), momenta[s].FloatAt(0))
		if d := math.Abs(e - e0); d > worst {
			worst = d
		}
	}
	if worst > 1e-3*math.Abs(e0) {
		t.Fatalf("pendulum energy drifted by %g (energy %g)", worst, e0)
	}
}

func TestIntegrateVerletErrors(t *testing.T) {
	accel := func(q *core.Array) (*core.Array, error) {
		return core.MulF(q, -1), nil
	}
	q0 := mustFloats(t, []float64{1})
	p0 := mustFloats(t, []float64{0})
	if _, _, err := IntegrateVerlet(accel, 0, 1, q0, p0, 0); err == nil {
		t.Fatal("zero steps: want an error")
	}
	if _, _, err := IntegrateVerlet(accel, 0, 1, q0, mustFloats(t, []float64{0, 1}), 5); err == nil {
		t.Fatal("mismatched vectors: want an error")
	}
	if _, _, err := IntegrateVerlet(accel, 0, 1, mustFloats(t, []float64{}), mustFloats(t, []float64{}), 5); err == nil {
		t.Fatal("empty state: want an error")
	}
	boom := func(*core.Array) (*core.Array, error) { return nil, base.Errf("accel failed") }
	if _, _, err := IntegrateVerlet(boom, 0, 1, q0, p0, 5); err == nil {
		t.Fatal("accel error: want an error")
	}
	wrong := func(*core.Array) (*core.Array, error) { return mustFloats(t, []float64{1, 1}), nil }
	if _, _, err := IntegrateVerlet(wrong, 0, 1, q0, p0, 5); err == nil {
		t.Fatal("wrong accel shape: want an error")
	}
}

// TestIntegrateVerletFloat32State pins the dtype contract: a float32
// state integrates (read per element, widened exactly) and lands where
// the float64 run lands; RawFloats is nil for float32, so the old
// code silently integrated zeros.
func TestIntegrateVerletFloat32State(t *testing.T) {
	accel := func(q *core.Array) (*core.Array, error) {
		return core.MulF(q, -1), nil
	}
	q0, err := core.FromFloat32s([]float32{1}, 1)
	if err != nil {
		t.Fatalf("FromFloat32s: %v", err)
	}
	p0, err := core.FromFloat32s([]float32{0}, 1)
	if err != nil {
		t.Fatalf("FromFloat32s: %v", err)
	}
	positions, momenta, err := IntegrateVerlet(accel, 0, 2*math.Pi, q0, p0, 400)
	if err != nil {
		t.Fatalf("IntegrateVerlet: %v", err)
	}
	if got := positions[400].FloatAt(0); math.Abs(got-1) > 1e-3 {
		t.Fatalf("float32 state closed at %.6g, want ≈ 1", got)
	}
	if got := momenta[0].FloatAt(0); got != 0 {
		t.Fatalf("initial momentum = %g, want 0", got)
	}
}

// TestIntegrateVerletIntStateErrors pins the refusal of int states:
// they have no place in a continuous integrator and must error rather
// than panic or read as zeros.
func TestIntegrateVerletIntStateErrors(t *testing.T) {
	accel := func(q *core.Array) (*core.Array, error) {
		return core.MulF(q, -1), nil
	}
	qi, err := core.FromInts([]int64{1}, 1)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	pi, err := core.FromInts([]int64{0}, 1)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	if _, _, err := IntegrateVerlet(accel, 0, 1, qi, pi, 5); err == nil {
		t.Fatal("expected an error for an int state")
	}
	if _, _, err := IntegrateVerlet(accel, 0, 1, mustFloats(t, []float64{1}), pi, 5); err == nil {
		t.Fatal("expected an error for an int momentum with a float position")
	}
	if _, _, err := IntegrateVerlet(accel, 0, 1, qi, mustFloats(t, []float64{0}), 5); err == nil {
		t.Fatal("expected an error for an int position with a float momentum")
	}
}