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

package linalg

import (
	"math"
	"slices"
	"sync"

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

// Rank-revealing QR with column pivoting (Businger and Golub). The
// plain Householder QR of `QR` answers a factorisation but hides the
// rank: an exactly or nearly dependent column leaves a rounding-level
// R diagonal and `LeastSquares` has to refuse the system outright.
// Pivoting sweeps the column of largest remaining norm to the front at
// every step, which forces the dependence into the trailing block: the
// |R| diagonal comes out non-increasing and its decay is the rank
// decision, the same tolerance convention `LeastSquares` applies to
// its unpivoted diagonal.

// RRQR returns the rank-revealing QR factorisation with column
// pivoting of an m×n matrix a with m ≥ n: A·P = Q·R, with Q m×m
// orthogonal, R m×n upper triangular, and perm the column permutation
// P, meaning the j-th column of A·P is column perm[j] of A. The rank
// is the count of |R[i,i]| above n·eps·max|R|, the tolerance
// `LeastSquares` applies; with column pivoting the diagonal decays
// monotonically, so the count is where the decay crosses the floor.
// Complex matrices, non-2-D inputs and underdetermined shapes are
// refused, matching the dense surface.
func RRQR(a *core.Array) (q, r *core.Array, perm []int, rank int, err error) {
	const name = "RRQR"
	if a.Dtype() == core.Complex {
		return nil, nil, nil, 0, base.Errf("%s: complex matrices are not supported", name)
	}
	if a.NDim() != 2 {
		return nil, nil, nil, 0, base.Errf("%s: needs a 2-D matrix, got shape %s", name, base.ShapeText(a.Shape()))
	}
	m, n := a.Shape()[0], a.Shape()[1]
	if m == 0 || n == 0 {
		return nil, nil, nil, 0, base.Errf("%s: zero-sized matrix, got shape %s", name, base.ShapeText(a.Shape()))
	}
	if m < n {
		return nil, nil, nil, 0, base.Errf("%s: needs m ≥ n, got shape %s", name, base.ShapeText(a.Shape()))
	}
	rMat, qMat, perm := rrqrFactor(denseFloats(a, m, n), m, n, true)
	return floatsToArray(qMat, []int{m, m}), floatsToArray(rMat, []int{m, n}), perm, rrqrRankOf(rMat, n), nil
}

// RRQRRank returns the numerical rank a pivoted factorisation of a
// reports under the package's tolerance convention: the count of
// |R[i,i]| above n·eps·max|R|. It is the rank `RRQR` returns, computed
// without the orthogonal factor for callers that only want the count.
func RRQRRank(a *core.Array) (int, error) {
	const name = "RRQRRank"
	if a.Dtype() == core.Complex {
		return 0, base.Errf("%s: complex matrices are not supported", name)
	}
	if a.NDim() != 2 {
		return 0, base.Errf("%s: needs a 2-D matrix, got shape %s", name, base.ShapeText(a.Shape()))
	}
	m, n := a.Shape()[0], a.Shape()[1]
	if m == 0 || n == 0 {
		return 0, base.Errf("%s: zero-sized matrix, got shape %s", name, base.ShapeText(a.Shape()))
	}
	if m < n {
		return 0, base.Errf("%s: needs m ≥ n, got shape %s", name, base.ShapeText(a.Shape()))
	}
	rMat, _, _ := rrqrFactor(denseFloats(a, m, n), m, n, false)
	return rrqrRankOf(rMat, n), nil
}

// rrqrRankOf counts the diagonal entries of the flat m×n R above
// n·eps·max|R|, the rank convention `LeastSquares` states for its
// unpivoted diagonal.
func rrqrRankOf(rMat []float64, n int) int {
	rMax := 0.0
	for i := range n {
		if v := math.Abs(rMat[i*n+i]); v > rMax {
			rMax = v
		}
	}
	rankTol := float64(n) * base.EpsF * rMax
	rank := 0
	for i := range n {
		if math.Abs(rMat[i*n+i]) > rankTol {
			rank++
		}
	}
	return rank
}

// rrqrColumnBlock is the number of R columns one blocked reflector pass
// covers, the blocking qrRaw applies to the same sweep and for the same
// reason: R is row-major, so a lone column's update walks a stride-n
// pattern that touches one cache line per element, where a block of
// rrqrColumnBlock columns spans one line of float64 and walks each row
// once.
const rrqrColumnBlock = 8

// rrqrMinWork is the element-touch floor below which one worker's share
// of a reflector application or of a column-norm scan falls back to the
// calling goroutine: below it the dispatch costs more than the work.
const rrqrMinWork = 2048

// rrqrFactor runs the pivoted Householder sweep on a flat m×n matrix.
// The column of largest remaining norm moves to the front at every
// step; norms are recomputed each step rather than downdated, so no
// cancellation drift and no squared-overflow window: the cost is one
// pass over the working block per step, the same order as the
// reflector sweep itself. The Householder reflectors and their
// application follow `qrRaw`: scale before squaring, one reflector per
// step, strictly in sweep order, so the serial result is the
// factorisation the parallel plain QR also answers.
func rrqrFactor(aMat []float64, m, n int, wantQ bool) (rMat, qMat []float64, perm []int) {
	rMat = make([]float64, m*n)
	copy(rMat, aMat)
	perm = make([]int, n)
	for i := range perm {
		perm[i] = i
	}
	if wantQ {
		qMat = eye(m)
	}
	col := make([]float64, m)
	norms := make([]float64, n)
	for k := range n {
		best := rrqrPivot(rMat, m, n, k, perm, norms)
		if best == 0 {
			// Every remaining column is exactly zero: the trailing
			// block of R stays zero and the sweep is done.
			break
		}
		lv := col[:m-k]
		for i := range lv {
			lv[i] = rMat[(k+i)*n+k]
		}
		hh := householderVectorInto(lv, lv)
		if hh.beta == 0 {
			continue
		}
		rrqrApply(rMat, m, n, k, hh, qMat, nil)
	}
	return rMat, qMat, perm
}

// rrqrPivot selects the column of largest remaining norm among columns
// k..n−1, swaps it into column k and records the swap in perm. The
// norms themselves are computed one worker per contiguous column slab,
// each column by the serial two-pass accumulation stridedNorm2 always
// ran, and the largest is then picked by a serial scan whose strict
// comparison keeps the first of equal maxima, so the tie rule is
// unchanged. A zero return means every remaining column is exactly zero.
//
// A slab is a run of consecutive columns, which is what makes the scan
// affordable at all: column j and column j+1 share seven of every eight
// cache lines their strided walks touch, so a slab re-reads them from
// the cache a lone column at a time cannot. The crew only forms when one
// worker's slab holds a full dispatch floor of entries.
func rrqrPivot(rMat []float64, m, n, k int, perm []int, norms []float64) float64 {
	cols, rows := n-k, m-k
	engine.ParallelMin(cols, max((rrqrMinWork+rows-1)/rows, 1), func(start, end int) {
		for j := start; j < end; j++ {
			norms[j] = stridedNorm2(rMat, k*n+k+j, n, rows)
		}
	})
	best, bestJ := 0.0, k
	for j := range cols {
		if nrm := norms[j]; nrm > best {
			best, bestJ = nrm, k+j
		}
	}
	if best == 0 {
		return 0
	}
	if bestJ != k {
		rrqrSwapCols(rMat, m, n, bestJ, k)
		perm[bestJ], perm[k] = perm[k], perm[bestJ]
	}
	return best
}

// rrqrApply applies the reflector to the R columns k..n−1 in column
// blocks and, when qMat is non-nil, to the m rows of the orthonormal
// accumulator, or when bv is non-nil to the tail of the right-hand side.
// The targets are disjoint buffers, so they share one crew; the R blocks
// and the accumulator rows take contiguous slices, and every dot keeps
// its own i-ascending chain with the subtraction repeating the same
// operand order, so no crew size and no block boundary moves a bit.
func rrqrApply(rMat []float64, m, n, k int, hh householder, qMat, bv []float64) {
	ln := m - k
	blocks := (n - k + rrqrColumnBlock - 1) / rrqrColumnBlock
	items := blocks
	if qMat != nil {
		items += m
	}
	if bv != nil {
		items++
	}
	weight := (n - k) + m
	w := min(weight*ln/rrqrMinWork+1, engine.WorkersFor(items))
	if w < 2 {
		rrqrApplyWorker(rMat, m, n, k, hh, qMat, bv, blocks, items, 0, 1)
		return
	}
	var wg sync.WaitGroup
	for gi := range w {
		wg.Go(func() {
			rrqrApplyWorker(rMat, m, n, k, hh, qMat, bv, blocks, items, gi, w)
		})
	}
	wg.Wait()
}

// rrqrApplyWorker runs crew member gi of nw over its share of the
// reflector's targets: a contiguous slab of the R column blocks, a
// contiguous slab of the accumulator rows, and the right-hand side.
func rrqrApplyWorker(rMat []float64, m, n, k int, hh householder, qMat, bv []float64, blocks, items, gi, nw int) {
	ln := m - k
	// Target layout in the item index space: the R blocks first, then the
	// accumulator rows when there is an accumulator, then the right-hand
	// side when there is one.
	qEnd := blocks
	if qMat != nil {
		qEnd += m
	}
	var d [rrqrColumnBlock]float64
	for it := gi * items / nw; it < (gi+1)*items/nw; it++ {
		switch {
		case it < blocks:
			j0 := k + it*rrqrColumnBlock
			j1 := min(j0+rrqrColumnBlock, n)
			wd := j1 - j0
			_ = d[wd-1] // bounds-check proof: the block never exceeds rrqrColumnBlock
			clear(d[:wd])
			for jj := range wd {
				jcol := j0 + jj
				for i := range ln {
					d[jj] += hh.v[i] * rMat[(k+i)*n+jcol]
				}
			}
			for j := range wd {
				d[j] = hh.beta * d[j]
			}
			for i := range ln {
				vi := hh.v[i]
				base := (k+i)*n + j0
				for j := range wd {
					rMat[base+j] -= vi * d[j]
				}
			}
		case it < qEnd:
			row := it - blocks
			base := row*m + k
			dot := 0.0
			for i := range ln {
				dot += hh.v[i] * qMat[base+i]
			}
			wq := hh.beta * dot
			for i := range ln {
				qMat[base+i] -= hh.v[i] * wq
			}
		default:
			dot := 0.0
			for i := range ln {
				dot += hh.v[i] * bv[k+i]
			}
			wb := hh.beta * dot
			for i := range ln {
				bv[k+i] -= hh.v[i] * wb
			}
		}
	}
}

// rrqrSwapCols exchanges columns j1 and j2 of a flat m×n matrix.
func rrqrSwapCols(rMat []float64, m, n, j1, j2 int) {
	for i := range m {
		rMat[i*n+j1], rMat[i*n+j2] = rMat[i*n+j2], rMat[i*n+j1]
	}
}

// stridedNorm2 returns the L2 norm of count entries walking stride
// step from start, summed relative to the largest magnitude so a
// column near 1e154 does not overflow on the way to its norm.
func stridedNorm2(vals []float64, start, step, count int) float64 {
	maxAbs := 0.0
	for i := range count {
		if a := math.Abs(vals[start+i*step]); a > maxAbs {
			maxAbs = a
		}
	}
	if maxAbs == 0 {
		return 0
	}
	s := 0.0
	for i := range count {
		t := vals[start+i*step] / maxAbs
		s += t * t
	}
	return maxAbs * math.Sqrt(s)
}

// rrqrValidPerm reports whether p is a permutation of 0..n−1, the
// property the returned permutation's contract rests on.
func rrqrValidPerm(p []int) bool {
	seen := make([]bool, len(p))
	for _, v := range p {
		if v < 0 || v >= len(p) || seen[v] {
			return false
		}
		seen[v] = true
	}
	return !slices.Contains(seen, false)
}

// SolveRRQR solves min ‖A·x − b‖₂ through the pivoted factorisation.
// On full rank the answer is the ordinary back-substitution through R
// against c = Qᵀb. On rank-deficient input it is the minimum-norm
// least-squares solution: the rank r is read off the R diagonal decay,
// and the leading r×n trapezoidal block R₁ answers its underdetermined
// system in minimum norm through x = R₁ᵀ(R₁R₁ᵀ)⁻¹c₁, which squares the
// condition of the leading block exactly the way the package's
// BᵀB-based SVD squares the spectrum: honest for a rank decision taken
// from a decay that has already been observed.
func SolveRRQR(a, b *core.Array) (*core.Array, error) {
	const name = "SolveRRQR"
	if a.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex matrices are not supported", name)
	}
	if b.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex right-hand sides are not supported", name)
	}
	if a.NDim() != 2 {
		return nil, base.Errf("%s: needs a 2-D matrix, got shape %s", name, base.ShapeText(a.Shape()))
	}
	m, n := a.Shape()[0], a.Shape()[1]
	if m == 0 || n == 0 {
		return nil, base.Errf("%s: zero-sized matrix, got shape %s", name, base.ShapeText(a.Shape()))
	}
	if m < n {
		return nil, base.Errf("%s: needs m ≥ n, got shape %s", name, base.ShapeText(a.Shape()))
	}
	if b.NDim() != 1 || b.Len() != m {
		return nil, base.Errf("%s: right-hand side must be a vector of length %d, got shape %s",
			name, m, base.ShapeText(b.Shape()))
	}
	// The sweep below carries b through the reflectors as they are
	// built, so c = Qᵀb falls out of the factorisation without the
	// orthogonal factor ever being materialised.
	rMat := make([]float64, m*n)
	copy(rMat, denseFloats(a, m, n))
	bv := vectorF64(b, m)
	perm := make([]int, n)
	for i := range perm {
		perm[i] = i
	}
	col := make([]float64, m)
	norms := make([]float64, n)
	for k := range n {
		if rrqrPivot(rMat, m, n, k, perm, norms) == 0 {
			break
		}
		lv := col[:m-k]
		for i := range lv {
			lv[i] = rMat[(k+i)*n+k]
		}
		hh := householderVectorInto(lv, lv)
		if hh.beta == 0 {
			continue
		}
		rrqrApply(rMat, m, n, k, hh, nil, bv)
	}
	x := make([]float64, n)
	switch rank := rrqrRankOf(rMat, n); {
	case rank == n:
		// Full rank: back-substitution through R.
		for i := n - 1; i >= 0; i-- {
			s := bv[i]
			for j := i + 1; j < n; j++ {
				s -= rMat[i*n+j] * x[j]
			}
			x[i] = s / rMat[i*n+i]
		}
	default:
		// Minimum norm through the leading trapezoidal block R₁:
		// R₁R₁ᵀw = c₁, x = R₁ᵀw. R₁R₁ᵀ is symmetric positive definite
		// at the rank the decay showed, and the dense `Solve` carries
		// the small r×r system.
		if rank > 0 {
			mm := make([]float64, rank*rank)
			for i := range rank {
				for j := range rank {
					s := 0.0
					for l := range n {
						s += rMat[i*n+l] * rMat[j*n+l]
					}
					mm[i*rank+j] = s
				}
			}
			w, err := Solve(floatsToArray(mm, []int{rank, rank}), floatsToArray(append([]float64(nil), bv[:rank]...), []int{rank}))
			if err != nil {
				return nil, base.Errf("%s: %w", name, err)
			}
			for i := range rank {
				for j := range n {
					x[j] += rMat[i*n+j] * w.FloatAt(i)
				}
			}
		}
	}
	// x lives in permuted coordinates; undo the permutation.
	out := make([]float64, n)
	for j := range n {
		out[perm[j]] = x[j]
	}
	return floatsToArray(out, []int{n}), nil
}