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

package integrate

import "sourcedock.dev/petrbalvin/tensor/internal/core"

import (
	"math"
	"testing"
)

// harmonicSystem returns f for y” = −y written as the first-order
// system u' = (v, −u).
func harmonicSystem(t float64, y *core.Array) (*core.Array, error) {
	return core.FromFloats([]float64{y.FloatAt(1), -y.FloatAt(0)}, 2)
}

// TestIntegrateBoundaryLinear shoots the classic y” = −y with
// y(0) = 0 and y(π/2) = 1: the free initial slope must come out as 1
// and the sampled trajectory must trace y = sin t.
func TestIntegrateBoundaryLinear(t *testing.T) {
	bc := BoundaryConditions{Start: []int{0}, End: []int{0}, EndValues: []float64{1}}
	times, states, err := IntegrateBoundary(harmonicSystem, 0, math.Pi/2,
		mustFloats(t, []float64{0, 0.5}), bc, 5, ODEOptions{RelTol: 1e-9, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBoundary: %v", err)
	}
	for i := range times {
		want := math.Sin(times[i])
		if math.Abs(states[i].FloatAt(0)-want) > 1e-6 {
			t.Fatalf("y(%.6g) = %.12g, want %.12g", times[i], states[i].FloatAt(0), want)
		}
	}
	if math.Abs(states[0].FloatAt(1)-1) > 1e-5 {
		t.Fatalf("shooting slope = %.12g, want 1", states[0].FloatAt(1))
	}
	if math.Abs(times[4]-math.Pi/2) > 1e-12 {
		t.Fatalf("last sample time %.14g, want π/2", times[4])
	}
}

// TestIntegrateBoundaryNonlinear shoots y” = 2y³ with y(0) = 1 and
// y(1) = 1/2, whose exact solution is y = 1/(t+1) with the initial
// slope −1, and checks the sampled path against it.
func TestIntegrateBoundaryNonlinear(t *testing.T) {
	f := func(t float64, y *core.Array) (*core.Array, error) {
		u := y.FloatAt(0)
		return core.FromFloats([]float64{y.FloatAt(1), 2 * u * u * u}, 2)
	}
	bc := BoundaryConditions{Start: []int{0}, End: []int{0}, EndValues: []float64{0.5}}
	times, states, err := IntegrateBoundary(f, 0, 1,
		mustFloats(t, []float64{1, -0.5}), bc, 5, ODEOptions{RelTol: 1e-10, AbsTol: 1e-13})
	if err != nil {
		t.Fatalf("IntegrateBoundary: %v", err)
	}
	for i := range times {
		want := 1 / (times[i] + 1)
		if math.Abs(states[i].FloatAt(0)-want) > 1e-6 {
			t.Fatalf("y(%.6g) = %.12g, want %.12g", times[i], states[i].FloatAt(0), want)
		}
	}
	if math.Abs(states[0].FloatAt(1)+1) > 1e-5 {
		t.Fatalf("shooting slope = %.12g, want −1", states[0].FloatAt(1))
	}
}

// TestIntegrateBoundaryVelocityEnd prescribes the velocity at t1
// instead of the position: y” = −y with y(0) = 1 and y'(1) = 0 has
// y = cos t + tan(1)·sin t, and the shooting unknown is the initial
// value of the very component the end condition watches.
func TestIntegrateBoundaryVelocityEnd(t *testing.T) {
	bc := BoundaryConditions{Start: []int{0}, End: []int{1}, EndValues: []float64{0}}
	times, states, err := IntegrateBoundary(harmonicSystem, 0, 1,
		mustFloats(t, []float64{1, 1}), bc, 4, ODEOptions{RelTol: 1e-9, AbsTol: 1e-12})
	if err != nil {
		t.Fatalf("IntegrateBoundary: %v", err)
	}
	slope := math.Tan(1)
	for i := range times {
		want := math.Cos(times[i]) + slope*math.Sin(times[i])
		if math.Abs(states[i].FloatAt(0)-want) > 1e-6 {
			t.Fatalf("y(%.6g) = %.12g, want %.12g", times[i], states[i].FloatAt(0), want)
		}
	}
}

// TestIntegrateBoundaryBlowUp pins the honest failure: y” = y³ with
// y(0) = 1 and y(1) = 100 demands a trajectory that reaches 100 only
// by skirting its own finite-time blow-up, so some trial integration
// fails and the shooting reports the error instead of an answer.
func TestIntegrateBoundaryBlowUp(t *testing.T) {
	f := func(t float64, y *core.Array) (*core.Array, error) {
		u := y.FloatAt(0)
		return core.FromFloats([]float64{y.FloatAt(1), u * u * u}, 2)
	}
	bc := BoundaryConditions{Start: []int{0}, End: []int{0}, EndValues: []float64{100}}
	if _, _, err := IntegrateBoundary(f, 0, 1, mustFloats(t, []float64{1, 0}), bc, 3,
		ODEOptions{MaxSteps: 2000}); err == nil {
		t.Fatal("expected the shooting to fail on the blow-up problem")
	}
}

// TestIntegrateBoundaryErrors pins the validation contract.
func TestIntegrateBoundaryErrors(t *testing.T) {
	y0 := mustFloats(t, []float64{1, 0})
	bc := BoundaryConditions{Start: []int{0}, End: []int{0}, EndValues: []float64{1}}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, y0, bc, 1, ODEOptions{}); err == nil {
		t.Fatal("expected an error for a single sample")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, y0,
		BoundaryConditions{Start: []int{0}}, 3, ODEOptions{}); err == nil {
		t.Fatal("expected an error when nothing is prescribed at t1")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, y0,
		BoundaryConditions{Start: []int{0, 1}, End: []int{0}, EndValues: []float64{1}},
		3, ODEOptions{}); err == nil {
		t.Fatal("expected an error for too many conditions")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, y0,
		BoundaryConditions{Start: []int{0}, End: []int{0}, EndValues: []float64{}},
		3, ODEOptions{}); err == nil {
		t.Fatal("expected an error when EndValues does not match End")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, y0,
		BoundaryConditions{Start: []int{5}, End: []int{1}, EndValues: []float64{0}},
		3, ODEOptions{}); err == nil {
		t.Fatal("expected an error for a Start index out of range")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, y0,
		BoundaryConditions{Start: []int{0}, End: []int{2}, EndValues: []float64{1}},
		3, ODEOptions{}); err == nil {
		t.Fatal("expected an error for an End index out of range")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, mustFloats(t, []float64{1, 0, 2}),
		BoundaryConditions{Start: []int{0, 0}, End: []int{2}, EndValues: []float64{1}},
		3, ODEOptions{}); err == nil {
		t.Fatal("expected an error for a repeated Start index")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, y0,
		BoundaryConditions{Start: []int{0}, End: []int{1, 1}, EndValues: []float64{0, 1}},
		3, ODEOptions{}); err == nil {
		t.Fatal("expected an error for a repeated End index")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, mustFloats(t, []float64{1}, 1),
		bc, 3, ODEOptions{}); err == nil {
		t.Fatal("expected an error for a rank-2 state")
	}
	if _, _, err := IntegrateBoundary(harmonicSystem, 0, 1, mustFloats(t, nil), bc, 3, ODEOptions{}); err == nil {
		t.Fatal("expected an error for an empty state")
	}
}