// 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"
)

// Direct complex singular value decomposition by Golub-Kahan
// bidiagonalisation. Reducing A through Aᴴ·A squares the condition
// number, so a tiny singular value of A carries the accuracy of the
// square of a tiny eigenvalue; the direct route never forms the
// Gramian. core.Complex Householder reflectors chosen with real beta leave
// a real bidiagonal matrix, and the Golub-Reinsch shifted QR iteration
// then works on real numbers only, folding its rotations back into the
// complex unitary factors.

// svdBidiagonalise reduces the complex m×n matrix a (m ≥ n) to upper
// bidiagonal form with real diagonal d and superdiagonal e,
// accumulating the left unitary into u (m×m) and the right one into v
// (n×n), both seeded with the identity. Real beta is the choice that
// keeps both bidiagonal bands real: the reflector maps its head to a
// real multiple of the first coordinate whatever the head's phase.
func svdBidiagonalise(a []complex128, m, n int) (d, e []float64, u, v []complex128) {
	u = make([]complex128, m*m)
	v = make([]complex128, n*n)
	for i := range m {
		u[i*m+i] = 1
	}
	for i := range n {
		v[i*n+i] = 1
	}
	d = make([]float64, n)
	e = make([]float64, n)
	// Reflector scratch reused across the sweep; each step uses the
	// prefix it needs.
	vecBuf := make([]complex128, m)
	vecRBuf := make([]complex128, n)
	for k := range n {
		// Left reflector on column k, rows k..m-1.
		norm := 0.0
		for i := k; i < m; i++ {
			norm += base.Real2(a[i*n+k])
		}
		norm = math.Sqrt(norm)
		if norm > 0 {
			beta := -norm
			if real(a[k*n+k]) < 0 {
				beta = norm
			}
			vh := 0.0
			vec := vecBuf[:m-k]
			for i := k; i < m; i++ {
				vec[i-k] = a[i*n+k]
				if i == k {
					vec[0] -= complex(beta, 0)
				}
				vh += base.Real2(vec[i-k])
			}
			if vh > 0 {
				vhx := complex(0, 0)
				for i := k; i < m; i++ {
					vhx += cmplx.Conj(vec[i-k]) * a[i*n+k]
				}
				tau := complex(1, 0) / vhx
				for j := k; j < n; j++ {
					s := complex(0, 0)
					for i := k; i < m; i++ {
						s += cmplx.Conj(vec[i-k]) * a[i*n+j]
					}
					s *= tau
					for i := k; i < m; i++ {
						a[i*n+j] -= s * vec[i-k]
					}
				}
				for r := range m {
					s := complex(0, 0)
					for j := k; j < m; j++ {
						s += u[r*m+j] * vec[j-k]
					}
					s *= cmplx.Conj(tau)
					for j := k; j < m; j++ {
						u[r*m+j] -= s * cmplx.Conj(vec[j-k])
					}
				}
			}
		}
		d[k] = real(a[k*n+k])
		a[k*n+k] = complex(d[k], 0)
		for i := k + 1; i < m; i++ {
			a[i*n+k] = 0
		}
		if k == n-1 {
			break
		}
		// Right reflector on row k, columns k+1..n-1.
		norm = 0.0
		for j := k + 1; j < n; j++ {
			norm += base.Real2(a[k*n+j])
		}
		norm = math.Sqrt(norm)
		if norm > 0 {
			beta := -norm
			if real(a[k*n+k+1]) < 0 {
				beta = norm
			}
			vh := 0.0
			vecR := vecRBuf[:n-k-1]
			for j := k + 1; j < n; j++ {
				vecR[j-k-1] = a[k*n+j]
				if j == k+1 {
					vecR[0] -= complex(beta, 0)
				}
				vh += base.Real2(vecR[j-k-1])
			}
			if vh > 0 {
				vhx := complex(0, 0)
				for j := k + 1; j < n; j++ {
					vhx += cmplx.Conj(vecR[j-k-1]) * a[k*n+j]
				}
				tau := complex(1, 0) / vhx
				for i := range m {
					s := complex(0, 0)
					for j := k + 1; j < n; j++ {
						s += a[i*n+j] * cmplx.Conj(vecR[j-k-1])
					}
					s *= tau
					for j := k + 1; j < n; j++ {
						a[i*n+j] -= s * vecR[j-k-1]
					}
				}
				for r := range n {
					s := complex(0, 0)
					for j := k + 1; j < n; j++ {
						s += v[r*n+j] * cmplx.Conj(vecR[j-k-1])
					}
					s *= tau
					for j := k + 1; j < n; j++ {
						v[r*n+j] -= s * vecR[j-k-1]
					}
				}
			}
			e[k] = real(a[k*n+k+1])
			a[k*n+k+1] = complex(e[k], 0)
			for j := k + 2; j < n; j++ {
				a[k*n+j] = 0
			}
		}
	}
	return d, e, u, v
}

// svdGolubReinsch diagonalises the real bidiagonal pair (d, e) in
// place, folding every rotation into the columns of u (m×m) and v
// (n×n). Each sweep is one implicit shifted QR step on BᵀB: a
// Wilkinson shift from the trailing two-by-two opens the sweep, and
// the bulge it raises is chased off the block's end by alternating
// left and right rotations. A zero leading diagonal deflates through a
// rotation chase that empties its column, which is how rank deficiency
// leaves the iteration honestly.
func svdGolubReinsch(d, e []float64, u, v []complex128, m, n int) error {
	const name = "SVDComplex"
	// Deflation floor relative to the bidiagonal norm, the same
	// backward-stability floor symmetricQr uses: a purely
	// neighbour-relative threshold never triggers when both diagonals
	// sit at rounding-zero, which the FMA contraction of the vector
	// build leaves the chase on, and the iteration then exhausts
	// itself over an already negligible block.
	scale := 0.0
	for i := range n {
		scale += d[i] * d[i]
		if i+1 < n {
			scale += 2 * e[i] * e[i]
		}
	}
	tolAbs := base.EpsF * math.Sqrt(scale)
	negligible := func(i int) bool {
		a := math.Abs(e[i])
		return a == 0 ||
			a <= base.EpsF*(math.Abs(d[i])+math.Abs(d[i+1])) ||
			a < tolAbs
	}
	for iter := 0; iter < 60*n+100; iter++ {
		// Retire negligible superdiagonal entries.
		for i := range n - 1 {
			if negligible(i) {
				e[i] = 0
			}
		}
		// Locate the trailing unreduced block [l..r].
		r := -1
		for i := n - 2; i >= 0; i-- {
			if e[i] != 0 {
				r = i + 1
				break
			}
		}
		if r < 0 {
			return nil // fully diagonal
		}
		l := r - 1
		for l > 0 && e[l-1] != 0 {
			l--
		}
		if d[l] == 0 || math.Abs(d[l]) < tolAbs {
			// The block's leading column is zero, or sits at the
			// rounding floor of the whole bidiagonal: the FMA
			// contraction of the vector build leaves a denormal there
			// where the portable build hits the exact zero, and a
			// shifted chase over such a block cycles without
			// deflating. Left rotations between row l and the rows
			// below push the superdiagonal e[l] through the block and
			// off it, deflating the negligible singular value without
			// touching the rest.
			w := e[l]
			for k := l + 1; k <= r; k++ {
				h := math.Hypot(w, d[k])
				c, s := d[k]/h, -w/h
				d[k] = h
				if k < r {
					w = s * e[k]
					e[k] *= c
				}
				for row := range m {
					ul, uk := u[row*m+l], u[row*m+k]
					u[row*m+l] = complex(c, 0)*ul + complex(s, 0)*uk
					u[row*m+k] = complex(-s, 0)*ul + complex(c, 0)*uk
				}
			}
			e[l] = 0
			continue
		}
		// Wilkinson shift from the trailing two-by-two of BᵀB.
		t11 := d[r-1] * d[r-1]
		if r-2 >= 0 {
			t11 += e[r-2] * e[r-2]
		}
		t22 := d[r]*d[r] + e[r-1]*e[r-1]
		t21 := d[r-1] * e[r-1]
		delta := (t11 - t22) / 2
		denom := math.Abs(delta) + math.Sqrt(delta*delta+t21*t21)
		mu := t22
		if denom > 0 {
			sign := 1.0
			if delta < 0 {
				sign = -1.0
			}
			mu = t22 - sign*t21*t21/denom
		}
		// Opening right rotation on columns (l, l+1), taken from the
		// first column of T − μI.
		g0 := d[l]*d[l] - mu
		h0 := math.Hypot(g0, d[l]*e[l])
		c0, s0 := g0/h0, d[l]*e[l]/h0
		dl, el, dl1 := d[l], e[l], d[l+1]
		d[l] = c0*dl + s0*el
		e[l] = -s0*dl + c0*el
		d[l+1] = c0 * dl1
		bulge := s0 * dl1 // at (l+1, l)
		for row := range n {
			vl, vl1 := v[row*n+l], v[row*n+l+1]
			v[row*n+l] = complex(c0, 0)*vl + complex(s0, 0)*vl1
			v[row*n+l+1] = complex(-s0, 0)*vl + complex(c0, 0)*vl1
		}
		// Chase: the left rotation kills the bulge and raises a
		// super-bulge; the right rotation kills that and re-raises the
		// bulge one step down, until it falls off the block's end.
		for k := l; k < r; k++ {
			// Left rotation on rows (k, k+1). A deflated diagonal meets
			// a deflated bulge with h = 0: the rotation is the identity
			// there, where the division would raise 0/0 NaNs that spill
			// through the factor.
			h := math.Hypot(d[k], bulge)
			c1, s1 := 1.0, 0.0
			if h > 0 {
				c1, s1 = d[k]/h, bulge/h
			}
			ekOld, dk1Old := e[k], d[k+1]
			d[k] = h
			e[k] = c1*ekOld + s1*dk1Old
			d[k+1] = -s1*ekOld + c1*dk1Old
			var sb float64
			if k < r-1 {
				sb = s1 * e[k+1]
				e[k+1] *= c1
			}
			for row := range m {
				uk, uk1 := u[row*m+k], u[row*m+k+1]
				u[row*m+k] = complex(c1, 0)*uk + complex(s1, 0)*uk1
				u[row*m+k+1] = complex(-s1, 0)*uk + complex(c1, 0)*uk1
			}
			if k >= r-1 {
				break
			}
			// Right rotation on columns (k+1, k+2).
			h2 := math.Hypot(e[k], sb)
			c2, s2 := e[k]/h2, sb/h2
			ek1Old, dk1Old, dk2Old := e[k+1], d[k+1], d[k+2]
			e[k] = h2
			d[k+1] = c2*dk1Old + s2*ek1Old
			e[k+1] = -s2*dk1Old + c2*ek1Old
			bulge = s2 * dk2Old
			d[k+2] = c2 * dk2Old
			for row := range n {
				vk1, vk2 := v[row*n+k+1], v[row*n+k+2]
				v[row*n+k+1] = complex(c2, 0)*vk1 + complex(s2, 0)*vk2
				v[row*n+k+2] = complex(-s2, 0)*vk1 + complex(c2, 0)*vk2
			}
		}
	}
	return base.Errf("SVDComplex: the bidiagonal iteration did not converge")
}