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

package linalg

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

// core.Complex decompositions. The real `Eigen` tridiagonalises a symmetric
// matrix with Householder reflections, a construction that has no
// direct complex analogue; `EigenComplex` instead runs the Jacobi
// iteration, which generalises cleanly to Hermitian matrices through a
// two-step rotation: a complex phase turn that makes the offending
// off-diagonal entry real, then a real plane rotation that zeroes it.
// Every step preserves the spectrum exactly and the iteration inherits
// the real Jacobi convergence argument, at the usual price of
// O(n³) per sweep.
//
// `SVDComplex` builds on `EigenComplex`: the singular vectors of A are
// the eigenvectors of the Hermitian A^H·A, the singular values are
// their square roots, and the left vectors follow as A·v/σ, completed
// through the null space by Gram-Schmidt. Squaring the spectrum costs
// accuracy on very small singular values, which the doc contract
// states plainly.

// EigenComplex returns the eigenvalues and eigenvectors of a complex
// Hermitian matrix. Values are real, sorted ascending; vectors has
// shape (n, n) with column j the unit eigenvector for values[j].
// The matrix must be square, complex and Hermitian within a scale-
// relative 1e-12 tolerance. Real symmetric matrices should use the
// faster `Eigen`.
func EigenComplex(a *core.Array) (values, vectors *core.Array, err error) {
	if a.Dtype() != core.Complex {
		return nil, nil, base.Errf("EigenComplex: needs a complex matrix, got dtype %s", a.Dtype())
	}
	if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] {
		return nil, nil, base.Errf("EigenComplex: 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("EigenComplex: zero-sized matrix, got shape %s", base.ShapeText(a.Shape()))
	}
	h := make([]complex128, n*n)
	if a.Dtype() == core.Complex && !a.Strided() && len(a.RawComplexes()) == n*n {
		copy(h, a.RawComplexes())
	} else {
		for i := range n * n {
			h[i] = a.ComplexAt(i)
		}
	}
	if !hermitianOK(h, n) {
		return nil, nil, base.Errf("EigenComplex: matrix is not Hermitian within 1e-12 tolerance")
	}
	// A non-finite entry slips through the Hermitian test above: the
	// difference of a NaN and its conjugate never exceeds the tolerance.
	// The sparse twin refuses it in its own guard, and so does this one,
	// before the Jacobi sweep burns its passes on a poisoned matrix.
	for i := range n {
		for j := range n {
			z := h[i*n+j]
			if math.IsNaN(real(z)) || math.IsNaN(imag(z)) || math.IsInf(real(z), 0) || math.IsInf(imag(z), 0) {
				return nil, nil, base.Errf("EigenComplex: entry [%d,%d] is not finite", i, j)
			}
		}
	}
	// The Jacobi sweep's convergence test is a sum of squared
	// magnitudes: outside the safe window it overflows to +Inf, which
	// makes every sweep look already converged, or underflows to 0,
	// which makes the threshold 0 and stops on the first check, and
	// either way the raw diagonal is returned as the spectrum. The
	// matrix is moved into the window and the values, which carry the
	// scale, are moved back before they are returned; the vectors are
	// scale-free.
	ws := windowScale(maxMagComplex(h))
	if ws != 1 {
		scaleComplexes(h, ws)
	}
	// The accumulated similarity: columns of v are eigenvectors.
	v := eyeComplex(n)
	if err := hermitianJacobi(h, v, n); err != nil {
		return nil, nil, base.Errf("EigenComplex: %w", err)
	}

	vals := make([]float64, n)
	for i := range n {
		vals[i] = real(h[i*n+i])
	}
	idx := sortAscIndices(vals)
	outVals := make([]float64, n)
	outVecs := make([]complex128, n*n)
	for j := range n {
		outVals[j] = vals[idx[j]]
		for i := range n {
			outVecs[i*n+j] = v[i*n+idx[j]]
		}
	}
	if ws != 1 {
		unscaleFloats(outVals, ws)
	}
	valuesArr := floatsToArray(outVals, []int{n})
	vecArr := core.New(core.Complex, []int{n, n}...)
	copy(vecArr.RawComplexes(), outVecs)
	return valuesArr, vecArr, nil
}

// hermitianJacobi diagonalises a Hermitian matrix in place, sweeping
// all pairs until the off-diagonal mass is at rounding level. The
// eigenvectors accumulate into v. An exhausted sweep without reaching
// the threshold is an error, never a silently unconverged spectrum.
func hermitianJacobi(h []complex128, v []complex128, n int) error {
	norm := 0.0
	for i := range n * n {
		norm += real(h[i])*real(h[i]) + imag(h[i])*imag(h[i])
	}
	norm = math.Sqrt(norm)
	// A purely relative convergence threshold: an absolute floor such
	// as max(1, norm) would declare a Hermitian matrix of norm below
	// the floor already diagonal and return its raw diagonal, identity
	// eigenvectors included.
	threshold := 1e-13 * norm
	// The per-entry skip must sit where a matrix whose every
	// off-diagonal entry is skippable also passes the aggregate
	// convergence test: with P pairs the off-norm of all-skipped entries
	// reaches skip·√P, so a skip at the plain threshold leaves a matrix
	// whose entries are uniformly just under it frozen above the test,
	// and the sweep exhausts its passes without turning a single
	// rotation.
	skip := threshold
	if pairs := float64(n * (n - 1) / 2); pairs > 1 {
		skip = threshold / math.Sqrt(pairs)
	}
	offMass := func() float64 {
		off := 0.0
		for p := range n {
			for q := p + 1; q < n; q++ {
				z := h[p*n+q]
				off += real(z)*real(z) + imag(z)*imag(z)
			}
		}
		return math.Sqrt(off)
	}
	for range 60 {
		if offMass() <= threshold {
			return nil
		}
		for p := range n {
			for q := p + 1; q < n; q++ {
				z := h[p*n+q]
				if base.AbsComplex(z) <= skip {
					continue
				}
				// Step 1: a diagonal phase turn makes the pair entry
				// real and positive.
				if imag(z) != 0 || real(z) < 0 {
					alpha := cmplx.Phase(z)
					rot := cmplx.Exp(complex(0, -alpha))
					// The similarity h = Dᴴ·h·D turns the pair entry
					// real; V = V·D scales the p-th component of every
					// vector.
					for j := range n {
						h[p*n+j] *= rot
					}
					for i := range n {
						h[i*n+p] *= cmplx.Conj(rot)
					}
					for i := range n {
						v[i*n+p] *= cmplx.Conj(rot)
					}
					z = h[p*n+q]
				}
				// Step 2: the real Jacobi rotation that zeroes the now
				// real entry, applied to rows and columns p, q.
				tau := (real(h[q*n+q]) - real(h[p*n+p])) / (2 * real(z))
				var t float64
				if tau >= 0 {
					t = 1 / (tau + math.Sqrt(1+tau*tau))
				} else {
					t = -1 / (-tau + math.Sqrt(1+tau*tau))
				}
				c := 1 / math.Sqrt(1+t*t)
				s := t * c
				cr, sr := complex(c, 0), complex(s, 0)
				// h = Gᴴ·h·G mixes rows and columns p, q; V = V·G
				// mixes the p, q components of every eigenvector.
				for j := range n {
					hp, hq := h[p*n+j], h[q*n+j]
					h[p*n+j] = cr*hp - sr*hq
					h[q*n+j] = sr*hp + cr*hq
				}
				for i := range n {
					hp, hq := h[i*n+p], h[i*n+q]
					h[i*n+p] = cr*hp - sr*hq
					h[i*n+q] = sr*hp + cr*hq
				}
				for i := range n {
					vp, vq := v[i*n+p], v[i*n+q]
					v[i*n+p] = cr*vp - sr*vq
					v[i*n+q] = sr*vp + cr*vq
				}
				h[p*n+q] = 0
				h[q*n+p] = 0
			}
		}
	}
	if offMass() <= threshold {
		return nil
	}
	return base.Errf("the Jacobi sweep failed to converge on a %d×%d Hermitian matrix in 60 passes", n, n)
}

// hermitianOK reports whether every entry satisfies a[i][j] equals
// conj(a[j][i]) within a purely relative tolerance: an absolute floor
// would admit an asymmetry that is a large fraction of a small-scale
// matrix, so the guard would approve a matrix that is not Hermitian
// whatever the size of its entries.
func hermitianOK(h []complex128, n int) bool {
	scale := 0.0
	for _, z := range h {
		if a := base.AbsComplex(z); a > scale {
			scale = a
		}
	}
	tol := 1e-12 * scale
	for i := range n {
		for j := i; j < n; j++ {
			z, m := h[i*n+j], h[j*n+i]
			if d := base.AbsComplex(z - cmplx.Conj(m)); d > tol {
				return false
			}
		}
	}
	return true
}

func eyeComplex(n int) []complex128 {
	out := make([]complex128, n*n)
	for i := range n {
		out[i*n+i] = 1
	}
	return out
}

// SVDComplex returns the thin singular value decomposition of a
// complex m×n matrix: A = U·diag(Σ)·Vᴴ with Σ sorted descending. U
// has shape (m, n), sigma (n,) and vᴴ (n, n), matching the real `SVD`'s
// shapes. Wide inputs (m < n) decompose the conjugate transpose and
// swap the factors. Because the singular values come from the
// eigenvalues of Aᴴ·A, their accuracy on very small singular values is
// that of a squared condition, fine for rank decisions, weaker for
// resolving near-null directions.
func SVDComplex(a *core.Array) (u, sigma, vH *core.Array, err error) {
	if a.Dtype() != core.Complex {
		return nil, nil, nil, base.Errf("SVDComplex: needs a complex matrix, got dtype %s", a.Dtype())
	}
	if a.NDim() != 2 {
		return nil, nil, nil, base.Errf("SVDComplex: needs a 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
	}
	m, n := a.Shape()[0], a.Shape()[1]
	if m == 0 || n == 0 {
		return nil, nil, nil, base.Errf("SVDComplex: zero-sized matrix, got shape %s", base.ShapeText(a.Shape()))
	}
	flat := make([]complex128, m*n)
	if a.Dtype() == core.Complex && !a.Strided() && len(a.RawComplexes()) == m*n {
		copy(flat, a.RawComplexes())
	} else {
		for i := range m * n {
			flat[i] = a.ComplexAt(i)
		}
	}
	// The bidiagonalisation's reflector norms are sums of raw squares:
	// outside the safe window a tiny matrix has norm 0, which skips the
	// reflector and silently zeroes the below-diagonal mass, and a huge
	// one overflows to +Inf, which puts NaN in beta and the update. The
	// matrix is moved into the window and the singular values, which
	// carry the scale, are moved back on return; the unitary factors are
	// scale-free.
	ws := windowScale(maxMagComplex(flat))
	if ws != 1 {
		scaleComplexes(flat, ws)
	}
	if m < n {
		// A = V₂·Σ·U₂ᴴ obtained from the tall decomposition of Aᴴ.
		// Shapes mirror the real `SVD`: U (m, m), Σ (m,), Vᴴ (m, n).
		ah := transposeConj(flat, m, n)
		u2, sigma2, v2, err := svdComplexTall(ah, n, m)
		if err != nil {
			return nil, nil, nil, err
		}
		unscaleFloats(sigma2, ws)
		uOut := core.New(core.Complex, []int{m, m}...)
		copy(uOut.RawComplexes(), v2)
		sigOut := floatsToArray(sigma2, []int{m})
		vhOut := core.New(core.Complex, []int{m, n}...)
		copy(vhOut.RawComplexes(), transposeConj(u2, n, m))
		return uOut, sigOut, vhOut, nil
	}
	uMat, sigVals, vMat, err := svdComplexTall(flat, m, n)
	if err != nil {
		return nil, nil, nil, err
	}
	unscaleFloats(sigVals, ws)
	uOut := core.New(core.Complex, []int{m, n}...)
	copy(uOut.RawComplexes(), uMat)
	sigOut := floatsToArray(sigVals, []int{n})
	vhOut := core.New(core.Complex, []int{n, n}...)
	copy(vhOut.RawComplexes(), transposeConj(vMat, n, n))
	return uOut, sigOut, vhOut, nil
}

// svdComplexTall decomposes the complex m×n matrix with m ≥ n by
// Golub-Kahan bidiagonalisation and the Golub-Reinsch shifted QR
// iteration on the real bidiagonal: A = U·diag(σ)·Vᴴ with σ sorted
// descending, U thin (m×n) and V square (n×n).
func svdComplexTall(flat []complex128, m, n int) ([]complex128, []float64, []complex128, error) {
	work := append([]complex128(nil), flat...)
	d, e, u, v := svdBidiagonalise(work, m, n)
	if err := svdGolubReinsch(d, e, u, v, m, n); err != nil {
		return nil, nil, nil, err
	}
	// Non-negative singular values, flipping the matching U column.
	for i := range n {
		if d[i] < 0 {
			d[i] = -d[i]
			for r := range m {
				u[r*m+i] = -u[r*m+i]
			}
		}
	}
	// Descending order, permuting the factor columns along.
	for i := range n {
		big := i
		for j := i + 1; j < n; j++ {
			if d[j] > d[big] {
				big = j
			}
		}
		if big != i {
			d[i], d[big] = d[big], d[i]
			for r := range m {
				u[r*m+i], u[r*m+big] = u[r*m+big], u[r*m+i]
			}
			for r := range n {
				v[r*n+i], v[r*n+big] = v[r*n+big], v[r*n+i]
			}
		}
	}
	thin := make([]complex128, m*n)
	for i := range m {
		copy(thin[i*n:(i+1)*n], u[i*m:i*m+n])
	}
	return thin, d, v, nil
}

func transposeConj(a []complex128, m, n int) []complex128 {
	out := make([]complex128, n*m)
	for i := range m {
		for j := range n {
			out[j*m+i] = cmplx.Conj(a[i*n+j])
		}
	}
	return out
}