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

package linalg

import (
	"math/rand/v2"
	"slices"
	"testing"

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

// The reference side of the minimum degree tests: the selection scan
// the frontier replaced, kept verbatim so the equivalence test can
// hold the frontier to the exact sequence of choices the scan makes,
// tie breaks included, and the ordering benchmark can put a number on
// the difference.

// minimumDegreeScan returns the elimination order the original scan
// selects: every step walks the surviving vertices, counts each one's
// remaining neighbours and keeps the first vertex of the smallest
// count, so ties break to the smallest index.
func minimumDegreeScan(c *SparseCSC) ([]int, error) {
	adj, err := symmetrisedAdjacency(c)
	if err != nil {
		return nil, err
	}
	// The absorption below merges the lists as plain ascending index
	// sets, so the breadth first search's (degree, index) order has to
	// go: re-sort by index, exactly as the production order does.
	for i := range adj {
		slices.Sort(adj[i])
	}
	n := c.Cols
	eliminated := make([]bool, n)
	order := make([]int, 0, n)
	for range n {
		p := -1
		best := 0
		for i := range n {
			if eliminated[i] {
				continue
			}
			d := 0
			for _, u := range adj[i] {
				if !eliminated[u] {
					d++
				}
			}
			if p == -1 || d < best {
				p = i
				best = d
			}
		}
		if p == -1 {
			return nil, base.Errf("minimumDegree: no vertex left to eliminate")
		}
		order = append(order, p)
		eliminated[p] = true
		absorbElementScan(adj, eliminated, p)
	}
	return order, nil
}

// absorbElementScan merges the element the eliminated vertex p leaves
// behind into every surviving neighbour's adjacency list, the scan's
// form of the absorption: adj(j) becomes adj(j) ∪ adj(p) \ {j}.
func absorbElementScan(adj [][]int, eliminated []bool, p int) {
	for _, j := range adj[p] {
		if eliminated[j] {
			continue
		}
		adj[j] = sortedSetUnionScan(adj[j], adj[p])
		if at, found := slices.BinarySearch(adj[j], j); found {
			adj[j] = slices.Delete(adj[j], at, at+1)
		}
	}
}

// sortedSetUnionScan merges two sorted unique slices into one sorted
// unique slice.
func sortedSetUnionScan(a, b []int) []int {
	out := make([]int, 0, len(a)+len(b))
	i, j := 0, 0
	for i < len(a) && j < len(b) {
		switch {
		case a[i] < b[j]:
			out = append(out, a[i])
			i++
		case b[j] < a[i]:
			out = append(out, b[j])
			j++
		default:
			out = append(out, a[i])
			i++
			j++
		}
	}
	out = append(out, a[i:]...)
	out = append(out, b[j:]...)
	return out
}

// patternCOO assembles a symmetric matrix from off-diagonal edges plus
// a diagonal heavy enough to keep the matrix positive definite, though
// the ordering reads the pattern alone.
func patternCOO(t *testing.T, n int, edges [][2]int) *core.SparseCOO {
	t.Helper()
	idx := make([]int64, 0, 2*len(edges)+2*n)
	vals := make([]float64, 0, 2*len(edges)+2*n)
	add := func(r, c int) {
		idx = append(idx, int64(r), int64(c))
		vals = append(vals, 1)
	}
	for _, e := range edges {
		add(e[0], e[1])
		add(e[1], e[0])
	}
	for i := range n {
		idx = append(idx, int64(i), int64(i))
		vals = append(vals, float64(n)+2)
	}
	indices, err := core.FromInts(idx, len(vals), 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	coo, err := core.NewSparseCOO(indices, floatsToArray(vals, []int{len(vals)}), []int{n, n})
	if err != nil {
		t.Fatalf("NewSparseCOO: %v", err)
	}
	return coo
}

// gridEdges returns the edges of the w×h grid graph.
func gridEdges(w, h int) [][2]int {
	at := func(x, y int) int { return y*w + x }
	edges := make([][2]int, 0, 2*w*h)
	for y := range h {
		for x := range w {
			if x+1 < w {
				edges = append(edges, [2]int{at(x, y), at(x+1, y)})
			}
			if y+1 < h {
				edges = append(edges, [2]int{at(x, y), at(x, y+1)})
			}
		}
	}
	return edges
}

// cubeEdges returns the edges of the w×h×d grid graph.
func cubeEdges(w, h, d int) [][2]int {
	at := func(x, y, z int) int { return (z*h+y)*w + x }
	edges := make([][2]int, 0, 3*w*h*d)
	for z := range d {
		for y := range h {
			for x := range w {
				if x+1 < w {
					edges = append(edges, [2]int{at(x, y, z), at(x+1, y, z)})
				}
				if y+1 < h {
					edges = append(edges, [2]int{at(x, y, z), at(x, y+1, z)})
				}
				if z+1 < d {
					edges = append(edges, [2]int{at(x, y, z), at(x, y, z+1)})
				}
			}
		}
	}
	return edges
}

// starEdges returns the edges of the star graph: a centre joined to
// every other vertex, the shape whose eliminations hand the centre its
// degree one neighbour at a time.
func starEdges(n int) [][2]int {
	edges := make([][2]int, 0, n-1)
	for i := 1; i < n; i++ {
		edges = append(edges, [2]int{0, i})
	}
	return edges
}

// completeEdges returns the edges of the complete graph on n vertices,
// the shape whose first elimination fills everything.
func completeEdges(n int) [][2]int {
	edges := make([][2]int, 0, n*(n-1)/2)
	for i := range n {
		for j := i + 1; j < n; j++ {
			edges = append(edges, [2]int{i, j})
		}
	}
	return edges
}

// relabelledEdges renames every vertex through a random permutation:
// the same pattern under an index order chosen to scatter the degree
// ties the tie break has to survive.
func relabelledEdges(rng *rand.Rand, edges [][2]int) [][2]int {
	label := make(map[int]int)
	for _, e := range edges {
		label[e[0]] = 0
		label[e[1]] = 0
	}
	names := make([]int, 0, len(label))
	for v := range label {
		names = append(names, v)
	}
	slices.Sort(names)
	order := rng.Perm(len(names))
	for i, v := range names {
		label[v] = order[i]
	}
	out := make([][2]int, len(edges))
	for i, e := range edges {
		out[i] = [2]int{label[e[0]], label[e[1]]}
	}
	return out
}

// randomEdges draws m distinct off-diagonal edges of an n-vertex
// graph, the irregular patterns the ordering exists for. A random
// clique rides along every few calls to force the fill the absorptions
// have to keep up with.
func randomEdges(rng *rand.Rand, n, m int) [][2]int {
	if n < 2 {
		return nil
	}
	largest := n * (n - 1) / 2
	m = min(m, largest)
	seen := make(map[[2]int]bool)
	edges := make([][2]int, 0, m)
	for len(edges) < m {
		i := rng.IntN(n)
		j := rng.IntN(n)
		if i == j {
			continue
		}
		e := [2]int{min(i, j), max(i, j)}
		if seen[e] {
			continue
		}
		seen[e] = true
		edges = append(edges, e)
	}
	if rng.IntN(4) == 0 {
		clique := min(n, 2+rng.IntN(n/4+1))
		start := rng.IntN(n - clique + 1)
		for i := start; i < start+clique; i++ {
			for j := i + 1; j < start+clique; j++ {
				e := [2]int{i, j}
				if !seen[e] {
					seen[e] = true
					edges = append(edges, e)
				}
			}
		}
	}
	return edges
}

// checkOrderMatchesScan factors the pattern's adjacency once for each
// side and requires the frontier order to equal the scan order entry
// for entry, and to be a permutation at all.
func checkOrderMatchesScan(t *testing.T, name string, coo *core.SparseCOO) {
	t.Helper()
	c, err := CSCFromCOO(coo)
	if err != nil {
		t.Fatalf("%s: CSCFromCOO: %v", name, err)
	}
	want, err := minimumDegreeScan(c)
	if err != nil {
		t.Fatalf("%s: scan: %v", name, err)
	}
	got, err := minimumDegree(c)
	if err != nil {
		t.Fatalf("%s: frontier: %v", name, err)
	}
	for i := range min(len(got), len(want)) {
		if got[i] != want[i] {
			t.Fatalf("%s: step %d eliminates %d, the scan eliminates %d", name, i, got[i], want[i])
		}
	}
	if len(got) != len(want) {
		t.Fatalf("%s: order lengths differ: %d vs %d", name, len(got), len(want))
	}
	sorted := slices.Clone(got)
	slices.Sort(sorted)
	for i := range sorted {
		if sorted[i] != i {
			t.Fatalf("%s: order entry %d holds %d; not a permutation", name, i, sorted[i])
		}
	}
}

// TestMinimumDegreeMatchesScan holds the frontier ordering to the
// reference scan: on every pattern below, both must return the same
// elimination order, tie breaks included. The permutation decides the
// factorisation's fill, so a single differing choice is a failure.
func TestMinimumDegreeMatchesScan(t *testing.T) {
	rng := rand.New(rand.NewPCG(2026, 9))
	structured := []struct {
		name  string
		n     int
		edges [][2]int
	}{
		{"empty", 12, nil},
		{"path-50", 50, gridEdges(50, 1)},
		{"star-50", 50, starEdges(50)},
		{"complete-30", 30, completeEdges(30)},
		{"grid-12x12", 144, gridEdges(12, 12)},
		{"grid-15x15-shuffled", 225, relabelledEdges(rng, gridEdges(15, 15))},
		{"grid-64x64", 4096, gridEdges(64, 64)},
		{"cube-6x6x6", 216, cubeEdges(6, 6, 6)},
		{"two-grids", 64 + 36, append(gridEdges(8, 8), func() [][2]int {
			shifted := gridEdges(6, 6)
			for i := range shifted {
				shifted[i][0] += 64
				shifted[i][1] += 64
			}
			return shifted
		}()...)},
	}
	for _, tc := range structured {
		checkOrderMatchesScan(t, tc.name, patternCOO(t, tc.n, tc.edges))
	}
	for range 400 {
		n := 1 + rng.IntN(90)
		edges := randomEdges(rng, n, rng.IntN(3*n+1))
		if rng.IntN(2) == 0 {
			edges = relabelledEdges(rng, edges)
		}
		checkOrderMatchesScan(t, "random", patternCOO(t, n, edges))
	}
}