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

package integrate

import (
	"math"
	"testing"

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

// Benchmarks for the package's heavy paths: the adaptive Dormand-Prince
// step loop, the stiff implicit schemes with their numerical
// Jacobians, the adaptive quadrature and cubature, and the PDE
// stencils.

// odeLinear builds the closed-form linear system y' = A·y with a
// stable diagonal A, the cheapest honest workload for an adaptive
// step loop, and returns f plus the analytic solution for callers
// that want it.
func odeLinear(n int) (func(float64, *core.Array) (*core.Array, error), []float64) {
	rates := make([]float64, n)
	for i := range rates {
		rates[i] = -0.25 * float64(i+1)
	}
	f := func(t float64, y *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		vals := out.RawFloats()
		ys := y.RawFloats()
		for i := range n {
			vals[i] = rates[i] * ys[i]
		}
		return out, nil
	}
	y0 := make([]float64, n)
	for i := range y0 {
		y0[i] = 1
	}
	return f, y0
}

func benchVector(b *testing.B, vals []float64) *core.Array {
	b.Helper()
	a, err := core.FromFloats(vals, len(vals))
	if err != nil {
		b.Fatal(err)
	}
	return a
}

func BenchmarkIntegrateODE(b *testing.B) {
	f, y0 := odeLinear(16)
	start := benchVector(b, y0)
	opts := ODEOptions{}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateODE(f, 0, 10, start, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateRK4(b *testing.B) {
	f, y0 := odeLinear(16)
	start := benchVector(b, y0)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateRK4(f, 0, 10, start, 2000); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateBackwardEuler(b *testing.B) {
	// A stiff diagonal system: rates from −1 to −1000.
	const n = 4
	f := func(t float64, y *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		vals := out.RawFloats()
		ys := y.RawFloats()
		for i := range n {
			vals[i] = -float64(i+1) * 100 * ys[i]
		}
		return out, nil
	}
	start := benchVector(b, []float64{1, 1, 1, 1})
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateBackwardEuler(f, 0, 1, start, 200, ODEOptions{}); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateBDF2(b *testing.B) {
	const n = 4
	f := func(t float64, y *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		vals := out.RawFloats()
		ys := y.RawFloats()
		for i := range n {
			vals[i] = -float64(i+1) * 100 * ys[i]
		}
		return out, nil
	}
	start := benchVector(b, []float64{1, 1, 1, 1})
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateBDF2(f, 0, 1, start, ODEOptions{}); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateDAE(b *testing.B) {
	// The linear index-1 circuit shape: one differential row, one
	// algebraic constraint, the Newton solve carrying the step.
	m, err := core.FromFloats([]float64{1, 0, 0, 0}, 2, 2)
	if err != nil {
		b.Fatal(err)
	}
	f := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-y.FloatAt(0), y.FloatAt(1) - y.FloatAt(0)}, 2)
	}
	start := benchVector(b, []float64{1, 1})
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateDAE(f, m, 0, 1, start, 200, DAEOptions{}); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateMidpoint(b *testing.B) {
	// The harmonic oscillator's quadratic H: the implicit stage is a
	// root find whose gradient is linear in z.
	const n = 8
	gradH := func(z *core.Array) (*core.Array, error) {
		out := core.New(core.Float, 2*n)
		vals := out.RawFloats()
		zs := z.RawFloats()
		for i := range n {
			vals[i] = zs[n+i]
			vals[n+i] = zs[i]
		}
		return out, nil
	}
	q0 := make([]float64, n)
	p0 := make([]float64, n)
	for i := range q0 {
		q0[i] = math.Sin(float64(i))
		p0[i] = math.Cos(float64(i))
	}
	qs := benchVector(b, q0)
	ps := benchVector(b, p0)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := IntegrateMidpoint(gradH, 0, 1, qs, ps, 50, MidpointOptions{}); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateFunction(b *testing.B) {
	f := func(x float64) (float64, error) { return math.Sin(x), nil }
	opts := QuadratureOptions{}
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := IntegrateFunction(f, 0, 100, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateND(b *testing.B) {
	f := func(x []float64) float64 {
		s := 0.0
		for _, v := range x {
			s += v * v
		}
		return math.Exp(-s)
	}
	lo := []float64{-2, -2, -2}
	hi := []float64{2, 2, 2}
	opts := CubatureOptions{Tolerance: 1e-6}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateND(f, lo, hi, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateHeat1D(b *testing.B) {
	n := 256
	u0 := make([]float64, n)
	for i := range u0 {
		u0[i] = math.Sin(float64(i) / float64(n) * math.Pi)
	}
	state := benchVector(b, u0)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateHeat1D(state, 1, 1.0/257, 0.1, 0.0002, 10, 0, 0); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateWave1D(b *testing.B) {
	n := 256
	u0 := make([]float64, n)
	v0 := make([]float64, n)
	for i := range u0 {
		u0[i] = math.Sin(float64(i) / float64(n) * math.Pi)
	}
	us := benchVector(b, u0)
	vs := benchVector(b, v0)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateWave1D(us, vs, 1, 1.0/257, 0.5, 0.002, 10); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateHeat2D(b *testing.B) {
	rows, cols := 32, 32
	u0 := make([]float64, rows*cols)
	for r := range rows {
		for c := range cols {
			u0[r*cols+c] = math.Sin(float64(c)/float64(cols)*math.Pi) * math.Sin(float64(r)/float64(rows)*math.Pi)
		}
	}
	state, err := core.FromFloats(u0, rows, cols)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateHeat2D(state, 1, 1.0/33, 1.0/33, 0.02, 0.0004, 5, 0, 0, 0, 0); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateWave2D(b *testing.B) {
	rows, cols := 32, 32
	u0 := make([]float64, rows*cols)
	v0 := make([]float64, rows*cols)
	for r := range rows {
		for c := range cols {
			u0[r*cols+c] = math.Sin(float64(c)/float64(cols)*math.Pi) * math.Sin(float64(r)/float64(rows)*math.Pi)
		}
	}
	us, err := core.FromFloats(u0, rows, cols)
	if err != nil {
		b.Fatal(err)
	}
	vs, err := core.FromFloats(v0, rows, cols)
	if err != nil {
		b.Fatal(err)
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateWave2D(us, vs, 1, 1.0/33, 1.0/33, 0.05, 0.002, 5); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkIntegrateVerlet(b *testing.B) {
	// Two coupled oscillators apiece: the acceleration reads the
	// neighbour spring terms.
	n := 32
	q0 := make([]float64, n)
	p0 := make([]float64, n)
	for i := range q0 {
		q0[i] = math.Sin(float64(i))
	}
	accel := func(q *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		vals := out.RawFloats()
		qs := q.RawFloats()
		for i := range n {
			l, r := 0.0, 0.0
			if i > 0 {
				l = qs[i-1]
			}
			if i < n-1 {
				r = qs[i+1]
			}
			vals[i] = l - 2*qs[i] + r
		}
		return out, nil
	}
	qs := benchVector(b, q0)
	ps := benchVector(b, p0)
	b.ReportAllocs()
	for b.Loop() {
		if _, _, err := IntegrateVerlet(accel, 0, 10, qs, ps, 500); err != nil {
			b.Fatal(err)
		}
	}
}