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

package linalg

import (
	"math"

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

// The general (non-Hermitian) sparse eigenproblem. Lanczos
// needs symmetry: its three-term recurrence exists because the Krylov
// space of a symmetric operator comes with an orthogonal basis for
// free. A general operator needs the full Arnoldi recurrence, one
// explicit orthogonalisation per column, and its projected matrix is
// upper Hessenberg rather than tridiagonal. Resonance problems in
// electromagnetics and damped oscillation systems produce exactly
// these operators, complex and non-Hermitian included.
//
// Like SpEigen, this is a one-shot Krylov method: the basis runs to
// min(n, max(2k, k+40)) columns, the small Hessenberg eigenproblem is
// solved exactly by the dense EigenGeneral, and the returned pairs are
// Ritz approximations whose accuracy improves with the budget. The
// restart-on-breakdown scheme mirrors lanczos: an exhausted direction
// closes its block with a zero coupling and a fresh orthogonal random
// direction opens the next.

// SpEigenGeneral returns the k eigenvalues of largest magnitude of a
// general real sparse matrix, each with its unit eigenvector: values
// is a complex128 vector ordered by descending magnitude (a real
// matrix may carry complex conjugate pairs), vectors a complex128
// (n, k) array whose column j belongs to values[j]. The matrix must be
// square and real; complex input belongs to SpEigenGeneralComplex,
// symmetric input gets a cheaper answer from SpEigen. gen seeds the
// start vector (nil uses a fixed seed). Like the dense EigenGeneral,
// and unlike the symmetry-checking SpEigen, entries are not screened
// for finiteness: a non-finite entry answers NaN Ritz pairs.
func SpEigenGeneral(s *core.SparseCOO, k int, gen *core.Generator) (values, vectors *core.Array, err error) {
	const name = "SpEigenGeneral"
	if s.Values.Dtype() == core.Complex {
		return nil, nil, base.Errf("%s: complex matrices belong to SpEigenGeneralComplex", name)
	}
	// The same square 2-D contract SpEigen enforces, and for the same
	// reason: the Krylov product indexes x by the stored column, so a
	// column beyond the row count reads out of range, and the read
	// happens inside a worker goroutine the caller cannot recover.
	// Ranks above two are refused too: the index walk would silently
	// ignore every dimension past the first two.
	if len(s.Shape) != 2 || s.Shape[0] != s.Shape[1] {
		return nil, nil, base.Errf("%s: needs a square 2-D sparse matrix, got shape %v", name, s.Shape)
	}
	c, err := cooToCSR(s, name)
	if err != nil {
		return nil, nil, err
	}
	return arnoldiEigen(c.matVec, c.n, k, gen, name, krylovF64)
}

// SpEigenGeneralComplex is SpEigenGeneral for complex128 sparse
// matrices: non-Hermitian operators included, which is the
// electromagnetics case.
func SpEigenGeneralComplex(s *core.SparseCOO, k int, gen *core.Generator) (values, vectors *core.Array, err error) {
	const name = "SpEigenGeneralComplex"
	c, err := cooToComplexCSR(s, name)
	if err != nil {
		return nil, nil, err
	}
	return arnoldiEigen(c.matVec, c.n, k, gen, name, krylovC128)
}

// arnoldiEigen runs the shared pipeline: Arnoldi basis, dense
// eigensolve of the projected Hessenberg, Ritz lift. The matvec
// closure hides whether the operator is real or complex, and kern
// supplies the element-type-specific primitives the basis arithmetic
// needs.
func arnoldiEigen[T scalar](matvec func(x, y []T), n, k int, gen *core.Generator, name string, kern krylovKernel[T]) (*core.Array, *core.Array, error) {
	if n == 0 {
		return nil, nil, base.Errf("%s: zero-sized matrix", name)
	}
	if k < 1 || k > n {
		return nil, nil, base.Errf("%s: k must be in [1, %d], got %d", name, n, k)
	}
	if gen == nil {
		gen = core.NewGenerator(spEigenSeed)
	}
	m := min(n, max(2*k, k+spEigenBlock))
	h, v := arnoldi(matvec, n, m, gen, kern)

	// The projected problem: the leading m×m block of the Hessenberg,
	// dense and small enough for the general eigensolver.
	hArr, err := zeros(core.Complex, []int{m, m})
	if err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	hc := hArr.RawComplexes()
	for j := range m {
		for i := 0; i <= j+1 && i < m; i++ {
			hc[i*m+j] = toComplex(h[i*m+j])
		}
	}
	vals, vecs, err := EigenGeneral(hArr)
	if err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	if k > vals.Len() {
		k = vals.Len()
	}
	outVals, err := zeros(core.Complex, []int{k})
	if err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	vc := outVals.RawComplexes()
	outVecs := make([]complex128, n*k)
	tmp := make([]complex128, n)
	for j := range k {
		vc[j] = vals.ComplexAt(j)
		// Ritz vector: lift the Hessenberg eigenvector through the
		// Arnoldi basis, u = V·y.
		clear(tmp)
		for p := range m {
			y := vecs.ComplexAt(p*m + j)
			if y == 0 {
				continue
			}
			kern.widenAxpy(tmp, v[p*n:(p+1)*n], y)
		}
		norm := 0.0
		for _, z := range tmp {
			norm += real(z)*real(z) + imag(z)*imag(z)
		}
		norm = math.Sqrt(norm)
		if norm > 0 {
			for i := range n {
				tmp[i] /= complex(norm, 0)
			}
		}
		for i := range n {
			outVecs[i*k+j] = tmp[i]
		}
	}
	ritzVecs, err := complexFromArray2D(outVecs, n, k)
	if err != nil {
		return nil, nil, base.Errf("%s: %w", name, err)
	}
	return outVals, ritzVecs, nil
}

// toComplex widens a scalar kernel value for the complex lift.
func toComplex[T scalar](v T) complex128 {
	switch x := any(v).(type) {
	case float64:
		return complex(x, 0)
	case complex128:
		return x
	}
	return 0
}

// arnoldi builds the Krylov basis v (columns of length n, m+1 of
// them) and the upper Hessenberg h ((m+1)×m, row-major) of the
// operator behind matvec. Full reorthogonalisation, twice per column,
// keeps the basis orthonormal to rounding level; a collapsed direction
// closes its block with a zero coupling and restarts in a fresh
// orthogonal random direction.
func arnoldi[T scalar](matvec func(x, y []T), n, m int, gen *core.Generator, kern krylovKernel[T]) (h, v []T) {
	// The basis never exceeds the space it spans: m above n would let
	// the truncation branch below hand back a shorter h than the m×m
	// read loop expects.
	if m > n {
		panic("arnoldi: the block size exceeds the matrix order")
	}
	h = make([]T, (m+1)*m)
	v = make([]T, n*(m+1))
	w := make([]T, n)
	scale := 0.0
	addScale := func(z T) {
		if a := absOf(z); a > scale {
			scale = a
		}
	}
	// normT is the one accumulation in the recurrence that squares its
	// inputs. A column whose entries are above about 1.3e154 overflows
	// the raw sum of squares to +Inf (the projected coupling becomes
	// infinite and the Ritz values come back wrong without a word), and
	// one below about 1.5e-162 underflows it to zero (the exhaustion
	// test reads that as a collapsed Krylov block and closes the block
	// on a direction the operator never exhausted). The sum therefore
	// runs raw while it lands in the normal range, which keeps the
	// ordinary-scale arithmetic bit-for-bit, and falls back to a
	// max-scaled accumulation outside it, the way norm2F64 and normC
	// do. NaN is carried through the raw path so a poisoned input stays
	// poisoned rather than reading as an exhausted block.
	normT := func(a []T) float64 {
		if s := absOf(kern.dot(a, a)); s <= math.MaxFloat64 && (s >= 0x1p-1022 || math.IsNaN(s)) {
			return math.Sqrt(s)
		}
		maxAbs := 0.0
		for _, z := range a {
			if m := absOf(z); m > maxAbs {
				maxAbs = m
			}
		}
		if maxAbs == 0 {
			return 0
		}
		s := 0.0
		for _, z := range a {
			r := absOf(z) / maxAbs
			s += r * r
		}
		return maxAbs * math.Sqrt(s)
	}
	start := func(col int) {
		q := v[col*n : (col+1)*n]
		kern.fillRand(q, gen)
		for range 2 {
			for p := range col {
				row := v[p*n : (p+1)*n]
				d := kern.dot(row, q)
				for i := range n {
					q[i] -= d * row[i]
				}
			}
		}
		if norm := normT(q); norm > 0 {
			kern.scaleInto(q, q, 1/norm)
		}
	}
	start(0)
	for j := range m {
		q := v[j*n : (j+1)*n]
		matvec(q, w)
		for i := 0; i <= j; i++ {
			row := v[i*n : (i+1)*n]
			d := kern.dot(row, w)
			h[i*m+j] = d
			for t := range n {
				w[t] -= d * row[t]
			}
		}
		// Full reorthogonalisation, twice: the first pass removes the
		// accumulated loss, the second what the first reintroduces.
		for range 2 {
			for i := 0; i <= j; i++ {
				row := v[i*n : (i+1)*n]
				d := kern.dot(row, w)
				h[i*m+j] += d
				for t := range n {
					w[t] -= d * row[t]
				}
			}
		}
		beta := normT(w)
		addScale(h[j*m+j])
		if j+1 < m {
			if beta <= float64(n)*base.EpsF*scale {
				// The Krylov space of this block is exhausted; close it
				// with a zero coupling and restart in a fresh
				// direction. m never exceeds n (asserted at the top),
				// so the basis always has room for the restart column.
				start(j + 1)
				continue
			}
			h[(j+1)*m+j] = realT[T](beta)
			kern.scaleInto(v[(j+1)*n:(j+2)*n], w, 1/beta)
		}
	}
	return h, v
}

// krylovKernel holds the element-type-specific primitives the shared
// Arnoldi recurrence calls once per vector pass: the conjugated inner
// product, the real scaling of a vector, the start-vector fill and the
// widening lift of a real basis row into the complex Ritz vector.
// Passing them in keeps one copy of the recurrence's arithmetic while
// the element type is dispatched on once per pass instead of once per
// element.
type krylovKernel[T scalar] struct {
	dot       func(a, b []T) T
	scaleInto func(dst, src []T, s float64)
	fillRand  func(dst []T, gen *core.Generator)
	widenAxpy func(dst []complex128, src []T, y complex128)
}

// krylovF64 is the kernel of a real operator.
var krylovF64 = krylovKernel[float64]{
	dot: dotF64,
	scaleInto: func(dst, src []float64, s float64) {
		for i := range dst {
			dst[i] = src[i] * s
		}
	},
	fillRand: func(dst []float64, gen *core.Generator) {
		for i := range dst {
			dst[i] = gen.NormalUnit()
		}
	},
	widenAxpy: func(dst []complex128, src []float64, y complex128) {
		for i, v := range src {
			dst[i] += y * complex(v, 0)
		}
	},
}

// krylovC128 is the kernel of a complex operator: the inner product
// conjugates its left operand, and the start vector draws a real and an
// imaginary part from the generator, in that order.
var krylovC128 = krylovKernel[complex128]{
	dot: dotC,
	scaleInto: func(dst, src []complex128, s float64) {
		for i := range dst {
			dst[i] = src[i] * complex(s, 0)
		}
	},
	fillRand: func(dst []complex128, gen *core.Generator) {
		for i := range dst {
			dst[i] = complex(gen.NormalUnit(), gen.NormalUnit())
		}
	},
	widenAxpy: func(dst []complex128, src []complex128, y complex128) {
		for i, v := range src {
			dst[i] += y * v
		}
	},
}