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

package linalg

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

import "math"

// Matrix exponential, evaluated by scaling and squaring over a
// diagonal Pade approximant (Higham, 2005). That is the construction
// LAPACK's *ex* routines and every mainstream runtime use, because it
// is backward stable and needs nothing but matrix products and one LU
// solve, both of which the library already has.
//
// The kernel is generic over the same scalar family the LU kernel runs
// on, so the complex path stays in complex128 end to end instead of
// splitting off into a parallel implementation.

// expmTheta holds the 1-norm bound below which each Pade degree keeps
// its approximation error inside double precision (Higham 2005,
// Table 10.1).
var expmTheta = map[int]float64{
	3:  1.495585217958292e-2,
	5:  2.539398330063230e-1,
	7:  9.504178996162932e-1,
	9:  2.097847961257068e+0,
	13: 5.371920351148152e+0,
}

// expmCoeffs returns the numerator coefficients b_j of the degree-m
// diagonal Pade approximant, as the integers of Higham's Table 10.1.
// The common factor across the row cancels in (V-U)^{-1}(V+U), so the
// unnormalised values are used directly, exactly as the reference
// implementations do.
func expmCoeffs(deg int) []float64 {
	switch deg {
	case 3:
		return []float64{120, 60, 12, 1}
	case 5:
		return []float64{30240, 15120, 3360, 420, 30, 1}
	case 7:
		return []float64{17297280, 8648640, 1995840, 277200, 25200, 1512, 56, 1}
	case 9:
		return []float64{17643225600, 8821612800, 2075673600, 302702400,
			30270240, 2162160, 110880, 3960, 90, 1}
	default:
		return []float64{64764752532480000, 32382376266240000, 7771770303897600,
			1187353796428800, 129060195264000, 10559470521600, 670442572800,
			33522128640, 1323241920, 40840800, 960960, 16380, 182, 1}
	}
}

// padeDegree picks the approximant degree and the number of squarings
// for a given 1-norm: the smallest degree whose bound covers the norm,
// else degree 13 with the matrix scaled down by a power of two.
func padeDegree(nA float64) (deg, s int) {
	for _, d := range []int{3, 5, 7, 9, 13} {
		if nA <= expmTheta[d] {
			return d, 0
		}
	}
	s = max(int(math.Ceil(math.Log2(nA/expmTheta[13]))), 0)
	return 13, s
}

// MatrixExp returns the matrix exponential exp(A) of a square matrix.
// Ints and float32 promote to float64; a complex matrix answers in
// complex128. A zero matrix gives the identity, and exp(A) for
// diagonal A is the exponential of the diagonal.
func MatrixExp(a *core.Array) (*core.Array, error) {
	if a.NDim() != 2 || a.Shape()[0] != a.Shape()[1] {
		return nil, base.Errf("MatrixExp: needs a square 2-D matrix, got shape %s", base.ShapeText(a.Shape()))
	}
	n := a.Shape()[0]
	if n == 0 {
		return nil, base.Errf("MatrixExp: zero-sized matrix, got shape %s", base.ShapeText(a.Shape()))
	}
	if a.Dtype() == core.Complex {
		m, err := squareComplexMatrix(a, "MatrixExp")
		if err != nil {
			return nil, err
		}
		r, err := expmKernel(m)
		if err != nil {
			return nil, err
		}
		flat := make([]complex128, n*n)
		for i := range n {
			copy(flat[i*n:(i+1)*n], r[i])
		}
		return core.FromComplexes(flat, n, n)
	}
	m, err := squareFloatMatrix(a, "MatrixExp")
	if err != nil {
		return nil, err
	}
	r, err := expmKernel(m)
	if err != nil {
		return nil, err
	}
	flat := make([]float64, n*n)
	for i := range n {
		copy(flat[i*n:(i+1)*n], r[i])
	}
	return floatsToArray(flat, []int{n, n}), nil
}

// expmKernel computes exp(a) by scaling and squaring over a Pade
// approximant. The approximant is r(A) = (V-U)^{-1}(V+U) where U holds
// the odd part and V the even part of the numerator polynomial, both
// evaluated by Horner in A^2; s squarings then undo the scaling.
func expmKernel[T scalar](a [][]T) ([][]T, error) {
	n := len(a)
	nA := matrixNorm1(a)
	if math.IsNaN(nA) || math.IsInf(nA, 0) {
		// The theta ladder divides by the norm and int(Ceil(...)) of a
		// non-finite value is implementation-defined; the Padé solve on
		// Inf data would return an all-NaN matrix with a nil error.
		return nil, base.Errf("MatrixExp: the matrix holds a non-finite entry (norm %g)", nA)
	}
	deg, s := padeDegree(nA)
	scaled := a
	if s > 0 {
		scaled = matScaleT(a, math.Ldexp(1, -s))
	}
	b := expmCoeffs(deg)
	a2 := matMulT(scaled, scaled)
	// The Horner sweeps alternate between two buffers instead of landing
	// a fresh matrix per coefficient: the product reads both operands
	// and writes the buffer the previous step left free, so the entries
	// the next multiply reads are the ones this one wrote.
	v, vFree := eyeScaled[T](n, b[deg-1]), newSquare[T](n)
	uIn, uFree := eyeScaled[T](n, b[deg]), newSquare[T](n)
	for j := deg - 3; j >= 0; j -= 2 {
		matMulAddEyeInto(vFree, a2, v, b[j])
		v, vFree = vFree, v
	}
	for j := deg - 2; j >= 1; j -= 2 {
		matMulAddEyeInto(uFree, a2, uIn, b[j])
		uIn, uFree = uFree, uIn
	}
	u := matMulT(scaled, uIn)
	r, err := solveRight("MatrixExp", matSubT(v, u), matAddT(v, u))
	if err != nil {
		return nil, err
	}
	if s > 0 {
		sq := newSquare[T](n)
		for range s {
			matMulTInto(sq, r, r)
			r, sq = sq, r
		}
	}
	return r, nil
}

// newSquare allocates an n×n matrix in the kernel's element type as n
// row views over one flat backing slice: two allocations instead of
// n+1. The rows stay disjoint, so in-place row swaps remain valid.
func newSquare[T scalar](n int) [][]T {
	back := make([]T, n*n)
	rows := make([][]T, n)
	for i := range n {
		rows[i] = back[i*n : (i+1)*n]
	}
	return rows
}

// realT widens a float64 into the kernel's element type, giving a zero
// imaginary part when the kernel runs on complex128.
func realT[T scalar](f float64) T {
	var zero T
	if _, ok := any(zero).(complex128); ok {
		return any(complex(f, 0)).(T)
	}
	return any(f).(T)
}

// eyeScaled returns c times the n×n identity.
func eyeScaled[T scalar](n int, c float64) [][]T {
	out := newSquare[T](n)
	cc := realT[T](c)
	for i := range n {
		out[i][i] = cc
	}
	return out
}

// matMulAddEyeInto writes a·m + c·I into out, which must be a square
// matrix of the same order aliasing neither operand.
func matMulAddEyeInto[T scalar](out, a, m [][]T, c float64) {
	matMulTInto(out, a, m)
	cc := realT[T](c)
	for i := range out {
		out[i][i] += cc
	}
}

// matMulT multiplies two square matrices in the kernel's element type
// and returns a fresh matrix.
func matMulT[T scalar](a, b [][]T) [][]T {
	out := newSquare[T](len(a))
	matMulTInto(out, a, b)
	return out
}

// matMulTInto writes the product a·b into out, which must alias neither
// operand. out arrives holding an earlier product, so it is cleared
// first: a fresh matrix is zero by construction and the accumulation
// below depends on that. The loop order keeps the inner product over p
// contiguous in both operands, matching the dense kernel's cache
// behaviour, and the zero multiplier is skipped the way the product has
// always skipped it, so a zero times a non-finite entry contributes
// nothing rather than a NaN.
func matMulTInto[T scalar](out, a, b [][]T) {
	n := len(a)
	var zero T
	for i := range n {
		oi := out[i]
		clear(oi)
		ai := a[i]
		for p := range n {
			aip := ai[p]
			if aip == zero {
				continue
			}
			bp := b[p]
			for j := range n {
				oi[j] += aip * bp[j]
			}
		}
	}
}

// matAddT returns the element-wise sum of two square matrices.
func matAddT[T scalar](a, b [][]T) [][]T {
	n := len(a)
	out := newSquare[T](n)
	for i := range n {
		oi, ai, bi := out[i], a[i], b[i]
		for j := range n {
			oi[j] = ai[j] + bi[j]
		}
	}
	return out
}

// matSubT returns the element-wise difference of two square matrices.
func matSubT[T scalar](a, b [][]T) [][]T {
	n := len(a)
	out := newSquare[T](n)
	for i := range n {
		oi, ai, bi := out[i], a[i], b[i]
		for j := range n {
			oi[j] = ai[j] - bi[j]
		}
	}
	return out
}

// matScaleT multiplies every element by the real scalar f.
func matScaleT[T scalar](a [][]T, f float64) [][]T {
	n := len(a)
	fT := realT[T](f)
	out := newSquare[T](n)
	for i := range n {
		oi, ai := out[i], a[i]
		for j := range n {
			oi[j] = ai[j] * fT
		}
	}
	return out
}

// matrixNorm1 returns the maximum absolute column sum, the norm the
// scaling step is measured against.
func matrixNorm1[T scalar](m [][]T) float64 {
	n := len(m)
	maxCol := 0.0
	for j := range n {
		sum := 0.0
		for i := range n {
			sum += absOf(m[i][j])
		}
		// A NaN column sum can never win the strict max below, so it
		// propagates explicitly: the caller reads the norm as the
		// finiteness screen.
		if math.IsNaN(sum) {
			return math.NaN()
		}
		if sum > maxCol {
			maxCol = sum
		}
	}
	return maxCol
}

// solveRight solves a·x = b for row-major square matrices and returns
// x row-major. a is consumed by the factorisation, so callers hand
// over a freshly built matrix; the column-oriented solveSystem is
// reused underneath rather than a second LU being written.
func solveRight[T scalar](name string, a, b [][]T) ([][]T, error) {
	n := len(a)
	colBack := make([]T, n*n)
	cols := make([][]T, n)
	for j := range n {
		col := colBack[j*n : (j+1)*n]
		for i := range n {
			col[i] = b[i][j]
		}
		cols[j] = col
	}
	sol, err := base.SolveSystem(name, a, cols)
	if err != nil {
		return nil, err
	}
	x := newSquare[T](n)
	for i := range n {
		for j := range n {
			x[i][j] = sol[j][i]
		}
	}
	return x, nil
}