// 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"
	"sourcedock.dev/petrbalvin/tensor/linalg"
)

// Benchmarks for the per-step scratch of the stiff solvers, the PDE
// stencil steps and the finite-element assemblies: the paths where
// allocation churn and repeated lookups, not the arithmetic, set the
// cost.

// perfVector wraps a fixed literal as a rank-1 array.
func perfVector(b *testing.B, vals []float64) *core.Array {
	b.Helper()
	a, err := core.FromFloats(vals, len(vals))
	if err != nil {
		b.Fatal(err)
	}
	return a
}

// perfStiffDecay builds the diagonal stiff system y' = −100(i+1)·y_i
// with every component started at one: the rates span three decades,
// so the step control stretches over the fast transient and the
// Jacobian stays diagonal and cheap to evaluate.
func perfStiffDecay(n int) func(t float64, y *core.Array) (*core.Array, error) {
	rates := make([]float64, n)
	for i := range rates {
		rates[i] = -100 * float64(i+1)
	}
	return 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
	}
}

// perfConstantState returns a vector of n ones.
func perfConstantState(n int) []float64 {
	vals := make([]float64, n)
	for i := range vals {
		vals[i] = 1
	}
	return vals
}

func BenchmarkROS4Stiff(b *testing.B) {
	const n = 32
	f := perfStiffDecay(n)
	start := perfVector(b, perfConstantState(n))
	opts := ODEOptions{RelTol: 1e-6, AbsTol: 1e-9}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateROS4(f, 0, 1, start, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkBDFVarStiff(b *testing.B) {
	const n = 32
	f := perfStiffDecay(n)
	start := perfVector(b, perfConstantState(n))
	opts := BDFVarOptions{RelTol: 1e-6, AbsTol: 1e-9}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateBDFVar(f, 0, 1, start, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkHeat1DStepLoop(b *testing.B) {
	const n = 256
	u0 := make([]float64, n)
	for i := range u0 {
		u0[i] = math.Sin(float64(i+1) / float64(n+1) * math.Pi)
	}
	state := perfVector(b, u0)
	// Two samples put every step inside the loop under test: the
	// published history costs one copy either way.
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateHeat1D(state, 1, 1.0/257, 0.05, 1e-4, 2, 0, 0); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkWave1DStepLoop(b *testing.B) {
	const n = 256
	u0 := make([]float64, n)
	v0 := make([]float64, n)
	for i := range u0 {
		u0[i] = math.Sin(float64(i+1) / float64(n+1) * math.Pi)
	}
	state := perfVector(b, u0)
	vel := perfVector(b, v0)
	b.ReportAllocs()
	for b.Loop() {
		if _, err := IntegrateWave1D(state, vel, 1, 1.0/257, 0.05, 1e-4, 2); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkHeat2DStepLoop(b *testing.B) {
	const rows, cols = 64, 64
	u0 := make([]float64, rows*cols)
	for r := range rows {
		for c := range cols {
			u0[r*cols+c] = math.Sin(float64(c+1)/float64(cols+1)*math.Pi) *
				math.Sin(float64(r+1)/float64(rows+1)*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/65, 1.0/65, 0.002, 2e-5, 2, 0, 0, 0, 0); err != nil {
			b.Fatal(err)
		}
	}
}

// perfSquareBoundary lists the boundary nodes of the m by m cell grid
// on the unit square: the bottom and top rows, then the interior
// nodes of the left and right columns.
func perfSquareBoundary(m int) []int {
	nodes := make([]int, 0, 4*m)
	for i := range m + 1 {
		nodes = append(nodes, i, m*(m+1)+i)
	}
	for j := 1; j < m; j++ {
		nodes = append(nodes, j*(m+1), j*(m+1)+m)
	}
	return nodes
}

func BenchmarkPoissonFEM2D(b *testing.B) {
	const m = 48
	mesh, err := GridTriangleMesh2D(0, 0, 1, 1, m, m)
	if err != nil {
		b.Fatal(err)
	}
	bound := perfSquareBoundary(m)
	values := make([]float64, len(bound))
	opts := FEMPoissonOptions{
		Kappa:           1,
		DirichletNodes:  bound,
		DirichletValues: values,
		Ordering:        linalg.SparseOrderingReverseCuthillMcKee,
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SolvePoissonFEM2D(mesh, nil, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// perfBoxBoundary lists the vertices of the box tetrahedral mesh that
// sit on the unit cube's surface.
func perfBoxBoundary(mesh *TetraMesh3D) []int {
	nodes := make([]int, 0, mesh.Vertices3())
	for i := range mesh.Vertices3() {
		x, y, z := mesh.Vertices[3*i], mesh.Vertices[3*i+1], mesh.Vertices[3*i+2]
		if x == 0 || x == 1 || y == 0 || y == 1 || z == 0 || z == 1 {
			nodes = append(nodes, i)
		}
	}
	return nodes
}

func BenchmarkPoissonFEM3D(b *testing.B) {
	const m = 8
	mesh, err := BoxTetraMesh3D(0, 0, 0, 1, 1, 1, m, m, m)
	if err != nil {
		b.Fatal(err)
	}
	bound := perfBoxBoundary(mesh)
	values := make([]float64, len(bound))
	opts := FEMPoisson3DOptions{
		Kappa:           1,
		DirichletNodes:  bound,
		DirichletValues: values,
		Ordering:        linalg.SparseOrderingReverseCuthillMcKee,
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SolvePoissonFEM3D(mesh, nil, opts); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkPoissonFEM3DLoad(b *testing.B) {
	const m = 5
	mesh, err := BoxTetraMesh3D(0, 0, 0, 1, 1, 1, m, m, m)
	if err != nil {
		b.Fatal(err)
	}
	bound := perfBoxBoundary(mesh)
	values := make([]float64, len(bound))
	src := func(x, y, z float64) float64 {
		return 3 * math.Pi * math.Pi * math.Sin(math.Pi*x) * math.Sin(math.Pi*y) * math.Sin(math.Pi*z)
	}
	opts := FEMPoisson3DOptions{
		Kappa:           1,
		DirichletNodes:  bound,
		DirichletValues: values,
		Ordering:        linalg.SparseOrderingReverseCuthillMcKee,
	}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := SolvePoissonFEM3D(mesh, src, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkIntegrateHeat2DBig is the same scheme on a grid large enough
// that the step sweeps have work to share: 512 lines of 512 unknowns per
// half-step.
func BenchmarkIntegrateHeat2DBig(b *testing.B) {
	const rows, cols = 512, 512
	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/513, 1.0/513, 0.02, 0.0004, 2, 0, 0, 0, 0); err != nil {
			b.Fatal(err)
		}
	}
}