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

// denseExpApply forms the whole exponential densely and multiplies it
// by v, the reference the Krylov projection is compared against.
func denseExpApply(t *testing.T, vals, v []float64, n int) []float64 {
	t.Helper()
	a, err := core.FromFloats(vals, n, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	ea, err := MatrixExp(a)
	if err != nil {
		t.Fatalf("MatrixExp: %v", err)
	}
	out := make([]float64, n)
	for i := range n {
		s := 0.0
		for j := range n {
			s += ea.FloatAt(i*n+j) * v[j]
		}
		out[i] = s
	}
	return out
}

// TestSpExpApplyMatchesDense compares the Krylov projection against
// the dense exponential on symmetric matrices, at both a full and a
// truncated Krylov dimension.
func TestSpExpApplyMatchesDense(t *testing.T) {
	cases := []struct {
		name  string
		n     int
		diag  float64
		off   float64
		steps int
	}{
		// n ≤ 40 gets the whole Krylov space, so the answer is exact.
		{name: "exact_small", n: 6, diag: 3, off: -1, steps: 0},
		{name: "exact_at_budget", n: 40, diag: 4, off: -1, steps: 0},
		// A truncated dimension on a larger matrix is approximate.
		{name: "truncated", n: 100, diag: 5, off: -1, steps: 40},
		// A near-diagonal matrix, where the action is close to
		// elementwise and the projection converges immediately.
		{name: "weakly_coupled", n: 30, diag: 2, off: -0.01, steps: 0},
	}
	for _, tt := range cases {
		t.Run(tt.name, func(t *testing.T) {
			vals := spdTridiagonal(tt.n, tt.diag, tt.off)
			sp := sparseFromDense(t, vals, tt.n)
			v := make([]float64, tt.n)
			for i := range tt.n {
				v[i] = math.Sin(float64(i+1)) * 0.5
			}
			vArr, err := core.FromFloats(v, tt.n)
			if err != nil {
				t.Fatalf("FromFloats v: %v", err)
			}
			got, err := SpExpApply(sp, vArr, tt.steps)
			if err != nil {
				t.Fatalf("SpExpApply: %v", err)
			}
			want := denseExpApply(t, vals, v, tt.n)
			for i := range tt.n {
				g, w := got.FloatAt(i), want[i]
				if math.Abs(g-w) > 1e-8*(1+math.Abs(w)) {
					t.Fatalf("element %d = %.12g, want %.12g", i, g, w)
				}
			}
		})
	}
}

// TestSpExpApplyDiagonal pins the closed form: for a diagonal matrix,
// exp(A)·v is the elementwise exponential of the diagonal times v.
func TestSpExpApplyDiagonal(t *testing.T) {
	const n = 5
	diag := []float64{1, 2, -1, 0.5, 3}
	vals := make([]float64, n*n)
	for i := range n {
		vals[i*n+i] = diag[i]
	}
	sp := sparseFromDense(t, vals, n)
	v := []float64{1, 1, 1, 1, 1}
	vArr, err := core.FromFloats(v, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	got, err := SpExpApply(sp, vArr, 0)
	if err != nil {
		t.Fatalf("SpExpApply: %v", err)
	}
	for i := range n {
		want := math.Exp(diag[i])
		if math.Abs(got.FloatAt(i)-want) > 1e-12*(1+math.Abs(want)) {
			t.Fatalf("element %d = %.12g, want %.12g", i, got.FloatAt(i), want)
		}
	}
}

// TestSpExpApplyZeroVector pins the degenerate contract: exp(A)·0 is
// the zero vector, returned without dividing by a zero norm.
func TestSpExpApplyZeroVector(t *testing.T) {
	const n = 4
	vals := spdTridiagonal(n, 3, -1)
	sp := sparseFromDense(t, vals, n)
	v, err := core.Zeros(core.Float, n)
	if err != nil {
		t.Fatalf("Zeros: %v", err)
	}
	got, err := SpExpApply(sp, v, 0)
	if err != nil {
		t.Fatalf("SpExpApply: %v", err)
	}
	for i := range n {
		if got.FloatAt(i) != 0 {
			t.Fatalf("element %d = %.12g, want 0 for a zero vector", i, got.FloatAt(i))
		}
	}
}

// TestSpExpApplyDeterminism checks reproducibility: the projection
// starts from v itself and draws nothing random.
func TestSpExpApplyDeterminism(t *testing.T) {
	const n = 50
	vals := spdTridiagonal(n, 4, -1)
	sp := sparseFromDense(t, vals, n)
	v := make([]float64, n)
	for i := range n {
		v[i] = float64(i+1) / float64(n)
	}
	vArr, err := core.FromFloats(v, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	g1, err := SpExpApply(sp, vArr, 0)
	if err != nil {
		t.Fatalf("SpExpApply #1: %v", err)
	}
	g2, err := SpExpApply(sp, vArr, 0)
	if err != nil {
		t.Fatalf("SpExpApply #2: %v", err)
	}
	for i := range n {
		if g1.FloatAt(i) != g2.FloatAt(i) {
			t.Fatalf("element %d = %.12g vs %.12g across runs", i, g1.FloatAt(i), g2.FloatAt(i))
		}
	}
}

// TestSpExpApplyStepsClamped checks that a step count past the
// dimension is clamped to it, rather than overrunning the space.
func TestSpExpApplyStepsClamped(t *testing.T) {
	const n = 4
	vals := spdTridiagonal(n, 3, -1)
	sp := sparseFromDense(t, vals, n)
	v := []float64{1, 2, 3, 4}
	vArr, err := core.FromFloats(v, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	got, err := SpExpApply(sp, vArr, 1000)
	if err != nil {
		t.Fatalf("SpExpApply: %v", err)
	}
	want := denseExpApply(t, vals, v, n)
	for i := range n {
		if math.Abs(got.FloatAt(i)-want[i]) > 1e-10*(1+math.Abs(want[i])) {
			t.Fatalf("element %d = %.12g, want %.12g", i, got.FloatAt(i), want[i])
		}
	}
}

// TestSpExpApplyRejectsInvalid pins the error contract for every input
// the projection cannot honestly answer.
func TestSpExpApplyRejectsInvalid(t *testing.T) {
	vector := func(n int) *core.Array {
		v, err := core.FromFloats(make([]float64, n), n)
		if err != nil {
			t.Fatalf("FromFloats: %v", err)
		}
		return v
	}
	t.Run("complex_sparse", func(t *testing.T) {
		idx, _ := core.FromInts([]int64{0, 0}, 1, 2)
		vals, _ := core.FromComplexes([]complex128{1}, 1)
		sp, err := core.NewSparseCOO(idx, vals, []int{1, 1})
		if err != nil {
			t.Fatalf("NewSparseCOO: %v", err)
		}
		if _, err := SpExpApply(sp, vector(1), 0); err == nil {
			t.Fatal("expected an error for a complex sparse matrix")
		}
	})
	t.Run("not_square", func(t *testing.T) {
		idx, _ := core.FromInts([]int64{0, 0}, 1, 2)
		vals, _ := core.FromFloats([]float64{1}, 1)
		sp, err := core.NewSparseCOO(idx, vals, []int{1, 2})
		if err != nil {
			t.Fatalf("NewSparseCOO: %v", err)
		}
		if _, err := SpExpApply(sp, vector(1), 0); err == nil {
			t.Fatal("expected an error for a non-square matrix")
		}
	})
	t.Run("zero_sized", func(t *testing.T) {
		idx, _ := core.FromInts(nil, 0, 2)
		vals, _ := core.FromFloats(nil, 0)
		sp := &core.SparseCOO{Indices: idx, Values: vals, Shape: []int{0, 0}}
		if _, err := SpExpApply(sp, vector(0), 0); err == nil {
			t.Fatal("expected an error for a zero-sized matrix")
		}
	})
	t.Run("vector_wrong_length", func(t *testing.T) {
		vals := spdTridiagonal(3, 3, -1)
		sp := sparseFromDense(t, vals, 3)
		if _, err := SpExpApply(sp, vector(2), 0); err == nil {
			t.Fatal("expected an error for a vector of the wrong length")
		}
	})
	t.Run("vector_not_rank1", func(t *testing.T) {
		vals := spdTridiagonal(2, 3, -1)
		sp := sparseFromDense(t, vals, 2)
		m, _ := core.FromFloats([]float64{1, 0, 0, 1}, 2, 2)
		if _, err := SpExpApply(sp, m, 0); err == nil {
			t.Fatal("expected an error for a rank-2 vector")
		}
	})
	t.Run("complex_vector", func(t *testing.T) {
		vals := spdTridiagonal(2, 3, -1)
		sp := sparseFromDense(t, vals, 2)
		v, _ := core.FromComplexes([]complex128{1, 1}, 2)
		if _, err := SpExpApply(sp, v, 0); err == nil {
			t.Fatal("expected an error for a complex vector")
		}
	})
	t.Run("asymmetric", func(t *testing.T) {
		sp := sparseFromDense(t, []float64{4, 1, 2, 4}, 2)
		if _, err := SpExpApply(sp, vector(2), 0); err == nil {
			t.Fatal("expected an error for an asymmetric matrix")
		}
	})
}