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

// Sparse matrix exponential applied to a vector. The dense
// `MatrixExp` forms the whole n×n exponential, which is O(n³) time and
// O(n²) memory and answers the question "what is exp(A)". Often the
// question is narrower: "what is exp(A)·v", for one vector, as arises
// in diffusion, network propagation and the action of a matrix
// function on an initial state. Forming the whole exponential to
// multiply one vector wastes both.
//
// SpExpApply computes exp(A)·v through a Krylov projection. The
// Krylov subspace K_m(A, v) = span{v, A·v, A²·v, …} is where the
// action of exp(A) on v actually lives, and on that subspace A
// restricts to a small tridiagonal matrix T. The answer is then
// Q·exp(T)·(‖v‖·e₁): one small dense exponential, which reuses the
// very `MatrixExp` already written, plus a lift back through the
// Lanczos basis.
//
// The subspace must be exactly K_m(A, v), so this variant of Lanczos
// stops when the recurrence collapses rather than restarting in a
// fresh direction the way the eigensolver does: a restarted direction
// would leave the Krylov space of v and the projection would no longer
// approximate exp(A)·v. A collapse simply means the space is smaller
// than the budget, which is exact rather than a failure.

// spExpSteps is the Krylov dimension used when the caller passes zero:
// 40, matching the block the sparse eigensolver budgets past its
// eigenpair count. A matrix of n ≤ 40 gets its whole space and an
// exact answer; above that the result is an approximation whose
// accuracy grows with the step count.
const spExpSteps = 40

// SpExpApply returns exp(A)·v for a real symmetric sparse A and a
// rank-1 vector v of length n, by Krylov projection onto the subspace
// generated by A and v.
//
// steps is the Krylov dimension: a value ≤ 0 uses min(n, 40), and a
// value ≥ n decomposes the whole space, so the answer is exact rather
// than approximate. The matrix must be square, real and symmetric,
// verified with the same 1e-12 tolerance the other sparse entry points
// apply. A zero v gives the zero vector exactly, since exp(A)·0 = 0.
func SpExpApply(a *core.SparseCOO, v *core.Array, steps int) (*core.Array, error) {
	if a.Values.Dtype() == core.Complex {
		return nil, base.Errf("SpExpApply: complex sparse matrices are not supported")
	}
	if len(a.Shape) != 2 || a.Shape[0] != a.Shape[1] {
		return nil, base.Errf("SpExpApply: needs a square 2-D sparse matrix, got shape %v", a.Shape)
	}
	n := a.Shape[0]
	if n == 0 {
		return nil, base.Errf("SpExpApply: zero-sized matrix, got shape %v", a.Shape)
	}
	if v.NDim() != 1 || v.Shape()[0] != n {
		return nil, base.Errf("SpExpApply: vector must have length %d, got shape %s", n, base.ShapeText(v.Shape()))
	}
	if v.Dtype() == core.Complex {
		return nil, base.Errf("SpExpApply: complex vectors are not supported")
	}
	c, err := symmetricCSR(a, "SpExpApply")
	if err != nil {
		return nil, err
	}
	if steps <= 0 {
		steps = min(n, spExpSteps)
	}
	if steps > n {
		steps = n
	}

	start := make([]float64, n)
	copy(start, contiguousF64(v))
	nrmv := norm2F64(start)
	if nrmv == 0 {
		return core.Zeros(core.Float, n)
	}
	for i := range n {
		start[i] /= nrmv
	}

	alphas, betas, basis := c.krylov(start, steps)

	// The small tridiagonal matrix that A restricts to on the Krylov
	// subspace; its exponential is the action in those coordinates.
	m := len(alphas)
	tMat := make([]float64, m*m)
	for i := range m {
		tMat[i*m+i] = alphas[i]
		if i+1 < m {
			tMat[i*m+i+1] = betas[i]
			tMat[(i+1)*m+i] = betas[i]
		}
	}
	tExp, err := MatrixExp(floatsToArray(tMat, []int{m, m}))
	if err != nil {
		return nil, err
	}

	// exp(T)·(‖v‖·e₁) is ‖v‖ times the first column of exp(T); the
	// lift back through the basis then needs no separate product.
	out := make([]float64, n)
	tc := contiguousF64(tExp)
	for p := range m {
		y := nrmv * tc[p*m]
		if y == 0 {
			continue
		}
		row := basis[p*n : (p+1)*n]
		for i := range n {
			out[i] += y * row[i]
		}
	}
	return floatsToArray(out, []int{n}), nil
}

// krylov runs the Lanczos recurrence from an explicit start vector,
// returning the diagonal alpha, the off-diagonal beta and the basis
// vectors flattened as n-wide rows. Unlike the eigensolver's variant
// it stops at a collapse rather than restarting, because the subspace
// has to stay K_m(A, v) for the projection to answer exp(A)·v; a
// smaller space reached before the budget is exact, not a failure.
func (c *sparseCSR) krylov(start []float64, steps int) (alphas, betas []float64, basis []float64) {
	n := c.n
	q := append([]float64(nil), start...)
	w := make([]float64, n)
	// The budget bounds the recurrence: one basis row and at most one
	// coefficient per step, so all three are sized once rather than grown.
	alphas = make([]float64, 0, steps)
	betas = make([]float64, 0, steps)
	basis = make([]float64, 0, steps*n)
	scale := 0.0
	addScale := func(v float64) {
		if a := math.Abs(v); a > scale {
			scale = a
		}
	}
	for range steps {
		basis = append(basis, q...)
		c.matVec(q, w)
		alpha := dotF64(q, w)
		alphas = append(alphas, alpha)
		addScale(alpha)
		// Strip the two previous directions out of w, the three-term
		// recurrence itself.
		for i := range n {
			w[i] -= alpha * q[i]
		}
		if len(alphas) > 1 {
			beta := betas[len(betas)-1]
			prev := basis[(len(alphas)-2)*n : (len(alphas)-1)*n]
			for i := range n {
				w[i] -= beta * prev[i]
			}
		}
		// Full reorthogonalisation, twice: one pass removes the
		// accumulated loss, the second what the first reintroduces.
		for range 2 {
			for p := range len(alphas) {
				row := basis[p*n : (p+1)*n]
				d := dotF64(w, row)
				for i := range n {
					w[i] -= d * row[i]
				}
			}
		}
		beta := norm2F64(w)
		// Purely scale-relative exhaustion threshold: an absolute floor
		// would truncate the Krylov space of a tiny-norm matrix at
		// dimension one and replace the first-order correction with its
		// projection, a relative error of order one.
		if beta <= float64(n)*base.EpsF*scale {
			// The Krylov space is exhausted; the vectors already
			// taken span the whole action of A on v.
			break
		}
		betas = append(betas, beta)
		addScale(beta)
		for i := range n {
			q[i] = w[i] / beta
		}
	}
	// The last beta couples to a vector that was never formed, so it
	// does not belong on the tridiagonal band.
	if len(betas) >= len(alphas) {
		betas = betas[:len(alphas)-1]
	}
	return alphas, betas, basis
}