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

package linalg

import (
	"testing"

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

// Benchmarks for the minimum degree ordering: the frontier selection
// against the reference scan beside it, on the mesh patterns the
// direct solvers are measured on. Both variants run in one binary, and
// the pair runs in both orders across the two parent benchmarks, so a
// drift of the machine between the two halves of a round cannot dress
// itself up as a difference between the variants. The small grid is
// there so a regression the large grids would drown stays visible.

// gridLaplacian3D builds the 7-point Laplacian on a w×h×d grid in
// row-major order: symmetric positive definite, and the 3-D mesh where
// an ordering's fill decisions cost the most.
func gridLaplacian3D(b *testing.B, w, h, d int) *core.SparseCOO {
	b.Helper()
	n := w * h * d
	idx := make([]int64, 0, 7*n)
	vals := make([]float64, 0, 7*n)
	add := func(r, c int, v float64) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, v)
	}
	at := func(x, y, z int) int { return (z*h+y)*w + x }
	for z := range d {
		for y := range h {
			for x := range w {
				add(at(x, y, z), at(x, y, z), 6)
				if x+1 < w {
					add(at(x, y, z), at(x+1, y, z), -1)
					add(at(x+1, y, z), at(x, y, z), -1)
				}
				if y+1 < h {
					add(at(x, y, z), at(x, y+1, z), -1)
					add(at(x, y+1, z), at(x, y, z), -1)
				}
				if z+1 < d {
					add(at(x, y, z), at(x, y, z+1), -1)
					add(at(x, y, z+1), at(x, y, z), -1)
				}
			}
		}
	}
	return sparseCOOFrom(b, n, idx, vals)
}

// orderingGrid is one mesh the ordering benchmarks run on.
type orderingGrid struct {
	name    string
	w, h, d int
}

// orderingGrids are the meshes: the small grid every regression would
// show on, the two 2-D meshes, and the 3-D mesh where the quadratic
// scan costs most.
var orderingGrids = []orderingGrid{
	{"2d-361", 19, 19, 1},
	{"2d-4096", 64, 64, 1},
	{"2d-16384", 128, 128, 1},
	{"3d-32768", 32, 32, 32},
}

// orderingInput builds one grid's matrix, with the CSC form the
// orderings consume beside the COO the constructor takes, outside
// every timed region.
func orderingInput(b *testing.B, g orderingGrid) (*core.SparseCOO, *SparseCSC) {
	b.Helper()
	var coo *core.SparseCOO
	if g.d == 1 {
		coo = gridLaplacian(b, g.w, g.h)
	} else {
		coo = gridLaplacian3D(b, g.w, g.h, g.d)
	}
	c, err := CSCFromCOO(coo)
	if err != nil {
		b.Fatal(err)
	}
	return coo, c
}

// orderingVariants are the two selections, named for the sub-benchmark
// that times them: the reference scan and the production frontier.
func orderingVariants() []struct {
	name string
	run  func(*SparseCSC) error
} {
	return []struct {
		name string
		run  func(*SparseCSC) error
	}{
		{"scan", func(c *SparseCSC) error {
			_, err := minimumDegreeScan(c)
			return err
		}},
		{"frontier", func(c *SparseCSC) error {
			_, err := minimumDegree(c)
			return err
		}},
	}
}

// orderingBattery runs the scan and the frontier against the same
// pattern, in the order the caller picks.
func orderingBattery(b *testing.B, reverse bool) {
	for _, g := range orderingGrids {
		_, c := orderingInput(b, g)
		variants := orderingVariants()
		if reverse {
			variants[0], variants[1] = variants[1], variants[0]
		}
		for _, v := range variants {
			variant := v
			b.Run(g.name+"/"+variant.name, func(b *testing.B) {
				b.ReportAllocs()
				for b.Loop() {
					if err := variant.run(c); err != nil {
						b.Fatal(err)
					}
				}
			})
		}
	}
}

// BenchmarkMinimumDegreeOrdering measures the ordering alone, scan
// first.
func BenchmarkMinimumDegreeOrdering(b *testing.B) {
	orderingBattery(b, false)
}

// BenchmarkMinimumDegreeOrderingRev measures the ordering alone,
// frontier first: the mirror of the other parent, so the pair's two
// halves alternate which variant pays for the position.
func BenchmarkMinimumDegreeOrderingRev(b *testing.B) {
	orderingBattery(b, true)
}

// BenchmarkSparseCholeskyMinimumDegreeFactor measures NewSparseCholesky
// end to end with the minimum degree ordering: the ordering sits inside
// the construction, so its cost is part of the number.
func BenchmarkSparseCholeskyMinimumDegreeFactor(b *testing.B) {
	for _, g := range orderingGrids {
		coo, _ := orderingInput(b, g)
		b.Run(g.name, func(b *testing.B) {
			b.ReportAllocs()
			for b.Loop() {
				if _, err := NewSparseCholesky(coo, SparseOrderingMinimumDegree); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}