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

package linalg

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

func TestMatrixSqrtLog(t *testing.T) {
	// A = [[2,1],[1,2]] has eigenvalues 3 and 1, both positive.
	a := mustFloats(t, []float64{2, 1, 1, 2}, 2, 2)
	sq, err := MatrixSqrt(a)
	if err != nil {
		t.Fatalf("MatrixSqrt: %v", err)
	}
	// √A·√A = A.
	prod, _ := core.MatMul2D(sq, sq)
	for i := range 4 {
		if math.Abs(prod.FloatAt(i)-a.FloatAt(i)) > 1e-10 {
			t.Fatalf("√A·√A[%d] = %v, want %v", i, prod.FloatAt(i), a.FloatAt(i))
		}
	}
	lg, err := MatrixLog(a)
	if err != nil {
		t.Fatalf("MatrixLog: %v", err)
	}
	// e^{ln A} = A: exponentiate the log element-wise and rebuild.
	// Simpler: check ln A against the eigen-decomposition directly.
	evals, evecs, _ := Eigen(a)
	lnA := make([]float64, 4)
	for i := range 2 {
		for j := range 2 {
			s := 0.0
			for k := range 2 {
				s += evecs.FloatAt(i*2+k) * math.Log(evals.FloatAt(k)) * evecs.FloatAt(j*2+k)
			}
			lnA[i*2+j] = s
		}
	}
	for i := range 4 {
		if math.Abs(lg.FloatAt(i)-lnA[i]) > 1e-10 {
			t.Fatalf("ln A[%d] = %v, want %v", i, lg.FloatAt(i), lnA[i])
		}
	}
	// Negative eigenvalue refused.
	neg := mustFloats(t, []float64{1, 2, 2, 1}, 2, 2)
	if _, err := MatrixSqrt(neg); err == nil {
		t.Fatal("expected an error for a non-PSD matrix")
	}
	if _, err := MatrixLog(neg); err == nil {
		t.Fatal("expected an error for a matrix with a non-positive eigenvalue")
	}
}

// TestMatrixFunctionsTinyScale pins the purely relative spectrum
// tolerance: a tiny positive-definite matrix must pass the
// positive-definiteness screen and an indefinite one of the same scale
// must be refused, where the old max(1, scale) floors did the
// opposite.
func TestMatrixFunctionsTinyScale(t *testing.T) {
	// MatrixLog of diag(1e-12, 2e-12): ln of the eigenvalues.
	tiny := mustFloats(t, []float64{1e-12, 0, 0, 2e-12}, 2, 2)
	lg, err := MatrixLog(tiny)
	if err != nil {
		t.Fatalf("MatrixLog(diag(1e-12, 2e-12)): %v", err)
	}
	want := []float64{math.Log(1e-12), math.Log(2e-12)}
	for i := range 2 {
		if math.Abs(lg.FloatAt(i*2+i)-want[i]) > 1e-9 {
			t.Fatalf("ln A[%d][%d] = %.16g, want %.16g", i, i, lg.FloatAt(i*2+i), want[i])
		}
	}
	for _, i := range []int{1, 2} {
		if math.Abs(lg.FloatAt(i)) > 1e-12 {
			t.Fatalf("ln A off-diagonal [%d] = %v, want 0", i, lg.FloatAt(i))
		}
	}
	// MatrixSqrt of the same matrix: √ of the eigenvalues.
	sq, err := MatrixSqrt(tiny)
	if err != nil {
		t.Fatalf("MatrixSqrt(diag(1e-12, 2e-12)): %v", err)
	}
	for i, w := range []float64{1e-6, math.Sqrt(2) * 1e-6} {
		if math.Abs(sq.FloatAt(i*2+i)-w) > 1e-15 {
			t.Fatalf("√A[%d][%d] = %.16g, want %.16g", i, i, sq.FloatAt(i*2+i), w)
		}
	}
	// An indefinite matrix at the same tiny scale is refused, not
	// silently clamped.
	indefinite := mustFloats(t, []float64{-1e-12, 0, 0, 2e-12}, 2, 2)
	if _, err := MatrixSqrt(indefinite); err == nil {
		t.Fatal("MatrixSqrt of a tiny indefinite matrix: want an error")
	}
	if _, err := MatrixLog(indefinite); err == nil {
		t.Fatal("MatrixLog of a tiny indefinite matrix: want an error")
	}
}

// TestMatrixLogComplex pins the complex route of the principal
// logarithm: a nonsingular complex matrix logs eigenvalue-wise, while
// a singular one is refused instead of walking into diverging-series
// garbage.
func TestMatrixLogComplex(t *testing.T) {
	diag := core.New(core.Complex, 2, 2)
	diag.RawComplexes()[0] = 3
	diag.RawComplexes()[3] = 4
	lg, err := MatrixLog(diag)
	if err != nil {
		t.Fatalf("MatrixLog(complex diag(3, 4)): %v", err)
	}
	want := []float64{math.Log(3), math.Log(4)}
	for i := range 2 {
		if v := lg.ComplexAt(i*2 + i); math.Abs(real(v)-want[i]) > 1e-10 || math.Abs(imag(v)) > 1e-10 {
			t.Fatalf("ln A[%d][%d] = %v, want ≈ %v", i, i, v, want[i])
		}
	}
	singular := core.New(core.Complex, 2, 2)
	singular.RawComplexes()[0] = 1
	if _, err := MatrixLog(singular); err == nil {
		t.Fatal("MatrixLog of a singular complex matrix: want an error")
	}
}