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

package linalg

import (
	"cmp"
	"slices"
	"sourcedock.dev/petrbalvin/tensor/internal/base"
)

// Fill-reducing orderings for the sparse factorisations. The
// symmetrised adjacency builder and the orderings live here; the
// factorisation consumes nothing but a permutation, so a new ordering
// slots in without touching the numeric code.

// symmetrisedAdjacency returns the adjacency lists of the pattern of
// A + Aᵀ with the diagonal dropped: for every stored entry (i, j),
// i ≠ j, i sits on j's list and j on i's. Every list is sorted by
// ascending degree and then index, which fixes the visiting order the
// breadth first search uses. The lists slice out of one flat arena
// sized by a counting pass, so the construction allocates twice
// whatever the vertex count; the degree snapshot counts the entries a
// list received, duplicates included, exactly as the incremental build
// measured them, and the degree-then-index sort is a total order, so
// the fill order inside a list cannot reach the result.
func symmetrisedAdjacency(c *SparseCSC) ([][]int, error) {
	adj, counts, err := symmetrisedAdjacencyBuild(c)
	if err != nil {
		return nil, err
	}
	for i := range adj {
		slices.SortStableFunc(adj[i], func(x, y int) int {
			if d := cmp.Compare(counts[x], counts[y]); d != 0 {
				return d
			}
			return cmp.Compare(x, y)
		})
	}
	return adj, nil
}

// adjacencyIndexOrdered returns the same adjacency lists in plain
// ascending index order, the order the minimum-degree absorption
// merges: it skips the degree-then-index sort the breadth first
// search needs and no index-ordered consumer does.
func adjacencyIndexOrdered(c *SparseCSC) ([][]int, error) {
	adj, _, err := symmetrisedAdjacencyBuild(c)
	return adj, err
}

// symmetrisedAdjacencyBuild builds the adjacency lists in ascending
// index order: one flat arena sized by a counting pass, every list
// sorted and compacted. The counts it also returns are the raw
// duplicate-inclusive neighbour counts the build measured, the values
// the degree-then-index ordering sorts by.
func symmetrisedAdjacencyBuild(c *SparseCSC) ([][]int, []int, error) {
	counts := make([]int, c.Cols)
	for j := range c.Cols {
		for p := c.ColStart[j]; p < c.ColStart[j+1]; p++ {
			i := c.RowIdx[p]
			if i == j {
				continue
			}
			if i < 0 || i >= c.Cols {
				return nil, nil, base.Errf("adjacency: row index %d out of range for %d columns", i, c.Cols)
			}
			counts[j]++
			counts[i]++
		}
	}
	total := 0
	for i := range c.Cols {
		total += counts[i]
	}
	arena := make([]int, total)
	adj := make([][]int, c.Cols)
	next := make([]int, c.Cols)
	off := 0
	for i := range c.Cols {
		adj[i] = arena[off : off+counts[i]]
		next[i] = off
		off += counts[i]
	}
	for j := range c.Cols {
		for p := c.ColStart[j]; p < c.ColStart[j+1]; p++ {
			i := c.RowIdx[p]
			if i == j {
				continue
			}
			arena[next[j]] = i
			next[j]++
			arena[next[i]] = j
			next[i]++
		}
	}
	for i := range c.Cols {
		slices.Sort(adj[i])
		adj[i] = slices.Compact(adj[i])
	}
	return adj, counts, nil
}

// breadthFirstSearch walks the component of start and returns the
// visit order, labelling level for every vertex it reaches. It
// initialises level to the unvisited state itself, so the caller can
// pass a scratch slice.
func breadthFirstSearch(adj [][]int, level []int, start int) []int {
	for i := range level {
		level[i] = -1
	}
	order := []int{start}
	level[start] = 0
	for head := 0; head < len(order); head++ {
		v := order[head]
		for _, u := range adj[v] {
			if level[u] < 0 {
				level[u] = level[v] + 1
				order = append(order, u)
			}
		}
	}
	return order
}

// deepestLevel returns the vertices of the largest level in visit
// order.
func deepestLevel(order []int, level []int) []int {
	max := 0
	for _, v := range order {
		if level[v] > max {
			max = level[v]
		}
	}
	last := make([]int, 0, len(order))
	for _, v := range order {
		if level[v] == max {
			last = append(last, v)
		}
	}
	return last
}

// pseudoPeripheralStart finds a start vertex whose eccentricity is
// close to the component's diameter: two sweeps of taking the
// shallowest-degree vertex of the deepest level, the standard
// pseudo-peripheral heuristic. Every choice is deterministic, so the
// ordering is a pure function of the pattern.
func pseudoPeripheralStart(adj [][]int, level []int, guess int) int {
	start := guess
	for range 2 {
		order := breadthFirstSearch(adj, level, start)
		last := deepestLevel(order, level)
		best := last[0]
		for _, v := range last[1:] {
			if len(adj[v]) < len(adj[best]) {
				best = v
			}
		}
		start = best
	}
	return start
}

// reverseCuthillMcKee returns the elimination order of the symmetric
// pattern: the vertices of every component in reverse breadth first
// order from a pseudo-peripheral start. The reverse of the search
// order is what shrinks the bandwidth, and with it the fill a
// triangular factorisation produces. Components are taken in
// ascending order of their smallest unreached vertex, so the result
// is a pure function of the pattern: position k holds the original
// index of the row eliminated k-th.
func reverseCuthillMcKee(c *SparseCSC) ([]int, error) {
	adj, err := symmetrisedAdjacency(c)
	if err != nil {
		return nil, err
	}
	level := make([]int, c.Cols)
	visited := make([]bool, c.Cols)
	order := make([]int, 0, c.Cols)
	for s := range c.Cols {
		if visited[s] {
			continue
		}
		start := pseudoPeripheralStart(adj, level, s)
		component := breadthFirstSearch(adj, level, start)
		for _, v := range component {
			visited[v] = true
		}
		for _, c := range slices.Backward(component) {
			order = append(order, c)
		}
	}
	return order, nil
}

// minimumDegree returns the elimination order of the symmetric
// pattern by minimum degree: at every step the uneliminated vertex
// with the fewest remaining neighbours is eliminated, its neighbours
// absorb its pattern (the element the elimination creates), and
// degree ties break to the smallest index, so the order is a pure
// function of the pattern. On irregular patterns it shrinks the fill
// well below what a bandwidth ordering reaches; on structured meshes
// the reverse Cuthill-McKee order is its match.
//
// The selection runs on a lazy frontier instead of a scan over the
// surviving vertices: a binary heap of (degree, vertex) pairs where a
// vertex enters when its degree is first known and its pair is rewritten
// in place whenever the degree moves, so the heap holds every surviving
// vertex exactly once and hands out the smallest (degree, index) among
// them, which is the first vertex of the smallest degree the scan would
// pick, at a logarithmic price per entry instead of a full sweep per
// elimination.
func minimumDegree(c *SparseCSC) ([]int, error) {
	adj, err := adjacencyIndexOrdered(c)
	if err != nil {
		return nil, err
	}
	n := c.Cols
	eliminated := make([]bool, n)
	degree := make([]int, n)
	frontier := &degreeFrontier{entries: make([]degreeEntry, 0, n)}
	arena := &intArena{}
	for i := range n {
		degree[i] = len(adj[i])
		frontier.push(degreeEntry{degree[i], i})
	}
	order := make([]int, 0, n)
	for range n {
		if frontier.len() == 0 {
			return nil, base.Errf("minimumDegree: no vertex left to eliminate")
		}
		p := frontier.pop().vertex
		order = append(order, p)
		eliminated[p] = true
		absorbElement(adj, degree, eliminated, p, frontier, arena)
	}
	return order, nil
}

// degreeEntry pairs a vertex with the degree it carried when it
// entered the frontier.
type degreeEntry struct {
	degree int
	vertex int
}

// before orders the pairs by degree, then index: the heap's top is
// the vertex the selection would take.
func (e degreeEntry) before(other degreeEntry) bool {
	if e.degree != other.degree {
		return e.degree < other.degree
	}
	return e.vertex < other.vertex
}

// degreeFrontier is a binary min-heap of degreeEntry pairs holding
// every vertex at most once: push rewrites the vertex's pair in place
// when one is already queued, so the heap never carries a stale pair
// and its top is always the vertex the selection takes. The absence of
// stale pairs keeps the heap at the vertex count instead of one entry
// per degree movement, which is what a sift-down per pop would
// otherwise charge.
type degreeFrontier struct {
	entries []degreeEntry
	pos     []int
}

// len reports how many pairs the frontier holds.
func (f *degreeFrontier) len() int { return len(f.entries) }

// push records the vertex's current degree: the pair of a vertex
// already queued moves to its new place, a vertex without a pair
// enters at the end. pos grows to cover the vertex and fills its new
// slots with -1, so an absent vertex reads -1 and a queued one reads
// its heap index.
func (f *degreeFrontier) push(e degreeEntry) {
	for len(f.pos) <= e.vertex {
		f.pos = append(f.pos, -1)
	}
	if at := f.pos[e.vertex]; at != -1 {
		f.entries[at] = e
		f.fix(at)
		return
	}
	f.pos[e.vertex] = len(f.entries)
	f.entries = append(f.entries, e)
	f.up(len(f.entries) - 1)
}

// pop removes and returns the smallest pair. It is the caller's job to
// check len first; the returned pair is never stale.
func (f *degreeFrontier) pop() degreeEntry {
	entries := f.entries
	top := entries[0]
	f.pos[top.vertex] = -1
	last := len(entries) - 1
	entries[0] = entries[last]
	f.entries = entries[:last]
	if last > 0 {
		f.pos[entries[0].vertex] = 0
		f.down(0)
	}
	return top
}

// fix restores the heap order around i after the pair at i changed.
func (f *degreeFrontier) fix(i int) {
	if !f.up(i) {
		f.down(i)
	}
}

// up lifts the pair at i until its parent stops outranking it, and
// reports whether it moved at all.
func (f *degreeFrontier) up(i int) bool {
	entries := f.entries
	moved := false
	for i > 0 {
		parent := (i - 1) / 2
		if !entries[i].before(entries[parent]) {
			break
		}
		entries[i], entries[parent] = entries[parent], entries[i]
		f.pos[entries[i].vertex] = i
		i = parent
		moved = true
	}
	f.pos[entries[i].vertex] = i
	return moved
}

// down drops the pair at i until both children stop outranking it.
func (f *degreeFrontier) down(i int) {
	entries := f.entries
	for {
		child := 2*i + 1
		if child >= len(entries) {
			break
		}
		if right := child + 1; right < len(entries) && entries[right].before(entries[child]) {
			child = right
		}
		if !entries[child].before(entries[i]) {
			break
		}
		entries[i], entries[child] = entries[child], entries[i]
		f.pos[entries[i].vertex] = i
		i = child
	}
	f.pos[entries[i].vertex] = i
}

// intArena hands out int slices from one bump-allocated backing block,
// so the minimum-degree absorption merges cost an allocation per block
// instead of one per merge. A returned slice owns its range: the merge
// appends into a zero-length view whose capacity is its exact upper
// bound, and the arena never hands the same range out twice. The
// capacity a merge leaves over returns to the bump space through
// release, which must reach the arena before its next alloc, while the
// released slice is still the most recent one. Backing blocks are kept
// alive by the adjacency lists that point into them, not by the arena.
type intArena struct {
	buf []int
	// bump is the offset of the next free entry in buf; the block is
	// replaced once the space behind bump no longer fits a request.
	bump int
}

// arenaBlockInts is the smallest backing block: big enough that the
// small early merges share one allocation, small enough that a block
// full of dead lists retires without pinning much memory.
const arenaBlockInts = 16384

// alloc returns a zero-length slice with capacity n. The arena works
// from an empty buffer as well as a partially consumed one: the bump
// offset starts at zero and the block check reads it against the block
// length, so a zero-value arena allocates rather than panics.
func (a *intArena) alloc(n int) []int {
	if len(a.buf)-a.bump < n {
		a.buf = make([]int, max(arenaBlockInts, n))
		a.bump = 0
	}
	out := a.buf[a.bump : a.bump : a.bump+n]
	a.bump += n
	return out
}

// release folds the capacity a merge left over back into the bump
// space, so the next alloc reuses it instead of moving past it. The
// slice must be the arena's most recent allocation and must not be
// released twice.
func (a *intArena) release(out []int) {
	a.bump -= cap(out) - len(out)
}

// absorbElement merges the element the eliminated vertex p leaves
// behind into every surviving neighbour's adjacency list, the fill
// the elimination creates: adj(j) becomes adj(j) ∪ adj(p) minus j and
// minus every eliminated vertex. The set difference matters: p sits on
// j's list and j on p's, so the plain union hands j its own index
// back, and a self entry both inflates the degree the selection reads
// and spreads to other lists through later absorptions. A neighbour's
// degree is the count of surviving entries in the merged list, the
// number a scan would recount at the same step, kept here where the
// list changes; a neighbour whose degree moved re-enters the frontier.
func absorbElement(adj [][]int, degree []int, eliminated []bool, p int, frontier *degreeFrontier, arena *intArena) {
	// The element list is compacted once, in place, before any merge:
	// p is eliminated and its list is dead, no vertex becomes
	// eliminated while the absorb runs, and every neighbour's merge
	// walks the same list, so one filter pass here replaces one per
	// merged element. The compaction is stable, so the merge order is
	// the list's own.
	b := adj[p]
	w := 0
	for _, u := range b {
		if !eliminated[u] {
			b[w] = u
			w++
		}
	}
	b = b[:w]
	for _, j := range b {
		before := degree[j]
		adj[j] = mergedElement(adj[j], b, j, eliminated, degree, arena)
		if degree[j] != before {
			frontier.push(degreeEntry{degree[j], j})
		}
	}
}

// mergedElement merges two sorted unique slices into one sorted unique
// slice without the self entry and without the eliminated entries, and
// stores the surviving count into degree[self]. The b side arrives
// already free of eliminated entries, so only the entries that come
// from a (and the merged equal pairs) pay the filter; b still carries
// self exactly once, which its own check drops. The degree the
// selection reads is the count of surviving entries either way, so the
// elimination order the frontier produces cannot move. The result lands
// in an arena slice capped at len(a)+len(b), which the merge cannot
// exceed: every entry of both lists is emitted at most once and the
// self entry never is.
func mergedElement(a, b []int, self int, eliminated []bool, degree []int, arena *intArena) []int {
	n := len(a) + len(b)
	out := arena.alloc(n)
	// The arena reserved exactly n entries, so the merge writes through
	// the stretched view and cuts the result back at the end.
	out = out[:n:n]
	w := 0
	i, j := 0, 0
	for i < len(a) && j < len(b) {
		u := a[i]
		switch v := b[j]; {
		case u < v:
			i++
			if u != self && !eliminated[u] {
				out[w] = u
				w++
			}
		case v < u:
			j++
			if v != self {
				out[w] = v
				w++
			}
		default:
			i++
			j++
			if u != self && !eliminated[u] {
				out[w] = u
				w++
			}
		}
	}
	for ; i < len(a); i++ {
		if u := a[i]; u != self && !eliminated[u] {
			out[w] = u
			w++
		}
	}
	for ; j < len(b); j++ {
		if v := b[j]; v != self {
			out[w] = v
			w++
		}
	}
	out = out[:w]
	degree[self] = w
	// The merge filled w of the capacity it reserved: the tail goes
	// back to the arena before the next merge carves past it.
	arena.release(out)
	return out
}