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

package linalg

import (
	"math"
	"math/cmplx"
	"slices"
	"sourcedock.dev/petrbalvin/tensor/internal/base"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// The general eigenproblem. Eigen answers real symmetric matrices and
// EigenComplex Hermitian ones, where structure pays for itself. A
// general square matrix has no structure to exploit, so EigenGeneral
// goes through the complex Schur form: the input is promoted to
// complex128, reduced to upper Hessenberg form by complex Householder
// reflectors, and driven to upper triangular form by the shifted QR
// iteration, one Wilkinson-shifted sweep of Givens rotations per step,
// deflating whenever a subdiagonal entry falls under a scale-relative
// tolerance. Working over the complex plane is what keeps the
// iteration single-shift: conjugate eigenvalue pairs are ordinary
// points there, where a real iteration would need the double-shift
// bulge chase to reach them.
//
// Eigenvectors come from the triangular Schur factor by back
// substitution and one multiplication with the accumulated unitary
// similarity, which keeps them orthonormal to rounding level even
// when the eigenvalues themselves are ill conditioned.

// EigenGeneral returns the eigenvalues and eigenvectors of a square
// matrix of any dtype; real and integer inputs are promoted to
// complex128. Values is a complex128 vector sorted descending by
// magnitude, bitwise ties broken by descending real part and then
// descending imaginary part; conjugate pairs share a magnitude only
// to rounding, so their relative order follows the rounding noise
// rather than the tiebreak. Vectors is a complex128 (n, n) array
// whose column j is the unit eigenvector for values[j]. Symmetric
// real matrices get a faster answer from Eigen and Hermitian matrices
// from EigenComplex.
func EigenGeneral(a *core.Array) (values, vectors *core.Array, err error) {
	if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] {
		return nil, nil, base.Errf("EigenGeneral: needs a square 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
	}
	n := a.Shape()[0]
	if n == 0 {
		return nil, nil, base.Errf("EigenGeneral: zero-sized matrix, got shape %s", base.ShapeText(a.Shape()))
	}
	h := make([]complex128, n*n)
	switch {
	case a.Dtype() == core.Complex && !a.Strided() && len(a.RawComplexes()) == n*n:
		copy(h, a.RawComplexes())
	case a.Dtype() == core.Float && !a.Strided() && len(a.RawFloats()) == n*n:
		for i := range n * n {
			h[i] = complex(a.RawFloats()[i], 0)
		}
	default:
		for i := range n * n {
			h[i] = a.ComplexAt(i)
		}
	}
	// Outside the safe window the reflector norms, the norm the shift
	// floors are taken against and the squared magnitudes the sweeps form
	// all leave the normal range: a tiny matrix collapses to zeros and a
	// huge one has vMax*vMax overflow, which zeroes beta so the
	// Hessenberg reduction is skipped and the QR iteration then fails on
	// a perfectly valid matrix. The matrix is moved into the window and
	// the eigenvalues, which carry the scale, are moved back on return;
	// the Schur vectors are scale-free.
	ws := windowScale(maxMagComplex(h))
	if ws != 1 {
		scaleComplexes(h, ws)
	}
	scale := 0.0
	for _, z := range h {
		if m := base.AbsComplex(z); m > scale {
			scale = m
		}
	}
	// q accumulates the unitary similarity: the input satisfies
	// A = q·H·qᴴ at every stage.
	q := eyeComplex(n)
	hessenbergComplex(h, q, n)
	if err := hessenbergQr(h, q, n, scale); err != nil {
		return nil, nil, err
	}

	vals := make([]complex128, n)
	for i := range n {
		vals[i] = h[i*n+i]
	}
	idx := sortMagDescIndices(vals)
	outVals := make([]complex128, n)
	for j := range n {
		outVals[j] = vals[idx[j]]
	}
	// Eigenvector j depends on the position of λ_j inside the Schur
	// factor, so they are built in place order and permuted afterwards.
	inPlace := make([]complex128, n*n)
	for j := range n {
		v := schurEigenVector(h, q, n, j, scale)
		for i := range n {
			inPlace[i*n+j] = v[i]
		}
	}
	outVecs := make([]complex128, n*n)
	for j := range n {
		for i := range n {
			outVecs[i*n+j] = inPlace[i*n+idx[j]]
		}
	}
	// Only the eigenvalues carry the matrix's scale.
	if ws != 1 {
		unscaleComplexes(outVals, ws)
	}
	valuesArr := core.New(core.Complex, []int{n}...)
	copy(valuesArr.RawComplexes(), outVals)
	vecArr := core.New(core.Complex, []int{n, n}...)
	copy(vecArr.RawComplexes(), outVecs)
	return valuesArr, vecArr, nil
}

// hessenbergComplex reduces h to upper Hessenberg form in place by
// complex Householder reflectors, accumulating the unitary similarity
// into q so that the original matrix satisfies A = q·H·qᴴ.
func hessenbergComplex(h, q []complex128, n int) {
	for k := range n - 2 {
		// The reflector zeroes column k below its subdiagonal; columns
		// already in form are skipped. Squared magnitudes are summed
		// relative to the column's largest entry so a column with
		// entries near 1e154 cannot overflow on the way to its norm.
		scale := 0.0
		for i := k + 1; i < n; i++ {
			if m := base.AbsComplex(h[i*n+k]); m > scale {
				scale = m
			}
		}
		below := 0.0
		if scale > 0 {
			for i := k + 2; i < n; i++ {
				m := base.AbsComplex(h[i*n+k]) / scale
				below += m * m
			}
		}
		if below == 0 {
			continue
		}
		x0 := h[(k+1)*n+k]
		norm := scale * math.Hypot(base.AbsComplex(x0)/scale, math.Sqrt(below))
		// alpha = −sign(x0)·‖x‖ lands the reflection on the far side
		// of x0, away from the cancellation zone.
		phase := complex(1, 0)
		if m := base.AbsComplex(x0); m > 0 {
			phase = x0 / complex(m, 0)
		}
		alpha := -phase * complex(norm, 0)
		v := make([]complex128, n-k-1)
		v[0] = x0 - alpha
		for i := 1; i < n-k-1; i++ {
			v[i] = h[(k+1+i)*n+k]
		}
		vMax := 0.0
		for _, z := range v {
			if m := base.AbsComplex(z); m > vMax {
				vMax = m
			}
		}
		vNorm2 := 0.0
		if vMax > 0 {
			for _, z := range v {
				re := real(z) / vMax
				im := imag(z) / vMax
				vNorm2 += re*re + im*im
			}
		}
		// beta = 2/(vᴴv) needs vMax*vMax to stay in the normal range. A
		// column whose entries are extreme in the matrix's own units
		// (far below it, say) overflows that square to +Inf, which makes
		// beta +0, or underflows it to 0, which makes beta +Inf and the
		// update NaN. Both are silent corruption of an otherwise valid
		// reduction, so the reflector is re-expressed in window units
		// with the scaling compensated in beta, exactly as
		// householderVectorInto does for the real reflectors.
		betaR := 2 / (vMax * vMax * vNorm2)
		if betaR <= 0 || math.IsInf(betaR, 0) {
			wsV := windowScale(vMax)
			for i := range v {
				v[i] = complex(real(v[i])*wsV, imag(v[i])*wsV)
			}
			w := vMax * wsV
			betaR = 2 / (w * w * vNorm2)
		}
		beta := complex(betaR, 0)

		// H = P·H with P = I − β·v·vᴴ, rows k+1..n−1, columns k..n−1.
		for j := k; j < n; j++ {
			s := complex(0, 0)
			for i := range v {
				s += cmplx.Conj(v[i]) * h[(k+1+i)*n+j]
			}
			s *= beta
			for i := range v {
				h[(k+1+i)*n+j] -= s * v[i]
			}
		}
		// H = H·P, every row, columns k+1..n−1.
		for i := range n {
			s := complex(0, 0)
			for j := range v {
				s += h[i*n+k+1+j] * v[j]
			}
			s *= beta
			for j := range v {
				h[i*n+k+1+j] -= s * cmplx.Conj(v[j])
			}
		}
		// q = q·P keeps the similarity: A = q·H·qᴴ.
		for i := range n {
			s := complex(0, 0)
			for j := range v {
				s += q[i*n+k+1+j] * v[j]
			}
			s *= beta
			for j := range v {
				q[i*n+k+1+j] -= s * cmplx.Conj(v[j])
			}
		}
		// Column k below the subdiagonal is zero by construction; make
		// it exactly so instead of carrying rounding residue.
		for i := k + 2; i < n; i++ {
			h[i*n+k] = 0
		}
	}
}

// hessenbergQr drives the shifted QR iteration on an upper Hessenberg
// matrix until every subdiagonal entry deflates, accumulating the
// Schur similarity into q. On return h is upper triangular: the Schur
// form of the original matrix, with its eigenvalues on the diagonal.
func hessenbergQr(h, q []complex128, n int, scale float64) error {
	// Purely relative to the matrix magnitude: an absolute floor would
	// treat a legitimate tiny-scale matrix (a 1e-20 rotation, say) as
	// one big deflated block and return zeros.
	floor := base.EpsF * scale
	negligible := func(i int) bool {
		local := base.EpsF * (base.AbsComplex(h[i*n+i]) + base.AbsComplex(h[(i-1)*n+i-1]))
		return base.AbsComplex(h[i*n+i-1]) <= math.Max(local, floor)
	}
	// Rotation scratch reused across sweeps; each sweep uses the first
	// hi−lo entries of each.
	cs := make([]complex128, n)
	sn := make([]complex128, n)
	hi := n - 1
	iter := 0
	for hi > 0 {
		for hi > 0 && negligible(hi) {
			h[hi*n+hi-1] = 0
			hi--
			iter = 0
		}
		if hi == 0 {
			return nil
		}
		lo := hi
		for lo > 0 && !negligible(lo) {
			lo--
		}
		if lo > 0 {
			// A negligible entry splits the matrix; the block above it
			// is settled in later rounds.
			h[lo*n+lo-1] = 0
		}
		iter++
		if iter > 100 {
			return base.Errf("EigenGeneral: QR iteration failed to converge at row %d", hi)
		}
		var mu complex128
		if iter%10 == 0 {
			// An exceptional shift breaks the rare cyclic pattern the
			// Wilkinson shift can settle into.
			mu = h[hi*n+hi] + complex(0.75*base.AbsComplex(h[hi*n+hi-1]), 0)
		} else {
			mu = wilkinsonShift(h, n, hi)
		}
		qrSweepComplex(h, q, n, lo, hi, mu, cs, sn)
	}
	return nil
}

// wilkinsonShift returns the eigenvalue of the trailing 2×2 block
// closest to its bottom-right entry: the shift that deflates the
// subdiagonal under it fastest.
func wilkinsonShift(h []complex128, n, i int) complex128 {
	a := h[(i-1)*n+i-1]
	b := h[(i-1)*n+i]
	c := h[i*n+i-1]
	d := h[i*n+i]
	delta := (a - d) / 2
	disc := cmplx.Sqrt(delta*delta + b*c)
	low := d + delta - disc
	high := d + delta + disc
	if base.AbsComplex(low-d) <= base.AbsComplex(high-d) {
		return low
	}
	return high
}

// qrSweepComplex performs one explicit single-shift QR step over the
// active block [lo, hi]: H = G·(H − μI)·Gᴴ + μI for the sequence of
// Givens rotations G that triangularises H − μI, with every rotation
// folded into q so the similarity stays exact. Each rotation satisfies
// G·[a; b] = [r; 0] with r real, which is what zeroes the subdiagonal
// one entry per step. cs and sn are the caller's scratch for the
// rotations; each sweep fully overwrites its first hi−lo entries.
func qrSweepComplex(h, q []complex128, n, lo, hi int, mu complex128, cs, sn []complex128) {
	for i := lo; i <= hi; i++ {
		h[i*n+i] -= mu
	}
	for j := lo; j < hi; j++ {
		a, b := h[j*n+j], h[(j+1)*n+j]
		r := math.Hypot(base.AbsComplex(a), base.AbsComplex(b))
		var c, s complex128
		if r == 0 {
			c, s = 1, 0
		} else {
			c, s = a/complex(r, 0), b/complex(r, 0)
		}
		cs[j-lo], sn[j-lo] = c, s
		// Left: rows j and j+1 over columns j..n−1.
		for k := j; k < n; k++ {
			x, y := h[j*n+k], h[(j+1)*n+k]
			h[j*n+k] = cmplx.Conj(c)*x + cmplx.Conj(s)*y
			h[(j+1)*n+k] = -s*x + c*y
		}
		h[(j+1)*n+j] = 0
	}
	for j := lo; j < hi; j++ {
		c, s := cs[j-lo], sn[j-lo]
		// Right: columns j and j+1 down to row j+1, the deepest row
		// the triangular factor can reach in either column.
		for i := 0; i <= j+1; i++ {
			x, y := h[i*n+j], h[i*n+j+1]
			h[i*n+j] = x*c + y*s
			h[i*n+j+1] = -x*cmplx.Conj(s) + y*cmplx.Conj(c)
		}
		for i := range n {
			x, y := q[i*n+j], q[i*n+j+1]
			q[i*n+j] = x*c + y*s
			q[i*n+j+1] = -x*cmplx.Conj(s) + y*cmplx.Conj(c)
		}
	}
	for i := lo; i <= hi; i++ {
		h[i*n+i] += mu
	}
}

// schurEigenVector builds the unit eigenvector for the diagonal entry
// j of a triangular Schur factor t with its similarity q (the input
// satisfies A = q·t·qᴴ): back substitution on the leading block gives
// the coordinates in Schur space, one multiplication with q lifts
// them back. A nearly multiple eigenvalue perturbs the denominator
// off exact zero, the standard defence against dividing by the
// spectrum's own degeneracy.
func schurEigenVector(t, q []complex128, n, j int, scale float64) []complex128 {
	x := make([]complex128, n)
	x[j] = 1
	lam := t[j*n+j]
	// Relative to the matrix magnitude for the same reason as the QR
	// floor above: tiny-scale spectra deserve their eigenvectors too.
	denFloor := base.EpsF * scale
	for i := j - 1; i >= 0; i-- {
		s := -t[i*n+j]
		for k := i + 1; k < j; k++ {
			s -= t[i*n+k] * x[k]
		}
		d := t[i*n+i] - lam
		if base.AbsComplex(d) < denFloor {
			d = complex(denFloor, 0)
		}
		if base.AbsComplex(d) == 0 {
			// The zero matrix: every denominator vanishes, and the
			// coordinate is unconstrained. Zero keeps the vector finite
			// where 0/0 would poison it with NaN.
			x[i] = 0
			continue
		}
		x[i] = s / d
	}
	v := make([]complex128, n)
	for i := range n {
		acc := complex(0, 0)
		for k := range j + 1 {
			acc += q[i*n+k] * x[k]
		}
		v[i] = acc
	}
	norm := 0.0
	for _, z := range v {
		norm += real(z)*real(z) + imag(z)*imag(z)
	}
	norm = math.Sqrt(norm)
	if norm > 0 {
		for i := range v {
			v[i] /= complex(norm, 0)
		}
	}
	return v
}

// sortMagDescIndices sorts indices so the values come descending by
// magnitude, bitwise ties broken by descending real part and then
// descending imaginary part. The sort is stable, so ties beyond the
// comparator keep their original order exactly as before.
func sortMagDescIndices(vals []complex128) []int {
	idx := make([]int, len(vals))
	for i := range idx {
		idx[i] = i
	}
	slices.SortStableFunc(idx, func(a, b int) int {
		ma, mb := base.AbsComplex(vals[a]), base.AbsComplex(vals[b])
		if ma != mb {
			if ma > mb {
				return -1
			}
			return 1
		}
		ra, rb := real(vals[a]), real(vals[b])
		if ra != rb {
			if ra > rb {
				return -1
			}
			return 1
		}
		ia, ib := imag(vals[a]), imag(vals[b])
		switch {
		case ia > ib:
			return -1
		case ia < ib:
			return 1
		default:
			return 0
		}
	})
	return idx
}