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

// pencilSample builds a deterministic symmetric a and a symmetric
// positive definite b of size n.
func pencilSample(n int) (*core.Array, *core.Array) {
	raw := make([]float64, n*n)
	for i := range n {
		for j := range n {
			raw[i*n+j] = math.Sin(float64(2*i + 3*j + 1))
		}
	}
	av := make([]float64, n*n)
	bv := make([]float64, n*n)
	for i := range n {
		for j := range n {
			av[i*n+j] = raw[i*n+j] + raw[j*n+i]
			bv[i*n+j] = math.Cos(float64(3*i+2*j)) + math.Cos(float64(3*j+2*i))
		}
		bv[i*n+i] += float64(n) + 1
	}
	a, _ := core.FromFloats(av, n, n)
	b, _ := core.FromFloats(bv, n, n)
	return a, b
}

// TestEigenGeneralisedDiagonal pins the pencil on a diagonal pair,
// where the eigenvalues are the entrywise ratios and the eigenvectors
// are the scaled unit vectors.
func TestEigenGeneralisedDiagonal(t *testing.T) {
	a := mustFloats(t, []float64{1, 0, 0, 4}, 2, 2)
	b := mustFloats(t, []float64{1, 0, 0, 2}, 2, 2)
	values, vectors, err := EigenGeneralised(a, b)
	if err != nil {
		t.Fatalf("EigenGeneralised: %v", err)
	}
	if math.Abs(values.FloatAt(0)-1) > 1e-12 || math.Abs(values.FloatAt(1)-2) > 1e-12 {
		t.Fatalf("values = (%.12g, %.12g), want (1, 2)", values.FloatAt(0), values.FloatAt(1))
	}
	// First eigenvector: e1 scaled to unit B norm = (1, 0); second:
	// e2 with 2·vᵀBv... v = (0, 1/√2) so vᵀBv = (1/2)·2 = 1.
	if math.Abs(vectors.FloatAt(0)-1) > 1e-12 || math.Abs(vectors.FloatAt(3)-1/math.Sqrt2) > 1e-12 {
		t.Fatalf("vectors = (%.12g, %.12g; %.12g, %.12g), want (1, 0; 0, 1/√2)",
			vectors.FloatAt(0), vectors.FloatAt(1), vectors.FloatAt(2), vectors.FloatAt(3))
	}
}

// TestEigenGeneralisedResidual checks the defining equations on a
// deterministic pencil: A·X = B·X·Λ to rounding, the vectors
// B-orthonormal, the values ascending.
func TestEigenGeneralisedResidual(t *testing.T) {
	n := 6
	a, b := pencilSample(n)
	values, vectors, err := EigenGeneralised(a, b)
	if err != nil {
		t.Fatalf("EigenGeneralised: %v", err)
	}
	scale := 0.0
	for i := range a.Len() {
		scale = math.Max(scale, math.Abs(a.FloatAt(i))+math.Abs(b.FloatAt(i)))
	}
	worst := 0.0
	for j := range n {
		lam := values.FloatAt(j)
		for i := range n {
			// (A·X − λ·B·X) at row i, column j.
			ax, bx := 0.0, 0.0
			for k := range n {
				ax += a.FloatAt(i*n+k) * vectors.FloatAt(k*n+j)
				bx += b.FloatAt(i*n+k) * vectors.FloatAt(k*n+j)
			}
			worst = math.Max(worst, math.Abs(ax-lam*bx))
		}
	}
	if worst > 1e-8*scale {
		t.Fatalf("pencil residual %.3g exceeds %.3g", worst, 1e-8*scale)
	}
	// B-orthonormality: Xᵀ·B·X = I.
	for j := range n {
		for i := j; i < n; i++ {
			s := 0.0
			for k := range n {
				for l := range n {
					s += vectors.FloatAt(k*n+i) * b.FloatAt(k*n+l) * vectors.FloatAt(l*n+j)
				}
			}
			want := 0.0
			if i == j {
				want = 1
			}
			if math.Abs(s-want) > 1e-8 {
				t.Fatalf("XᵀBX[%d][%d] = %.12g, want %.12g", i, j, s, want)
			}
		}
	}
	for i := 1; i < n; i++ {
		if values.FloatAt(i) < values.FloatAt(i-1) {
			t.Fatalf("values not ascending at %d", i)
		}
	}
}

// TestEigenGeneralisedErrors pins the validation contract.
func TestEigenGeneralisedErrors(t *testing.T) {
	a, b := pencilSample(3)
	cx, _ := core.FromComplexes([]complex128{1, 0, 0, 1}, 2, 2)
	if _, _, err := EigenGeneralised(cx, b); err == nil {
		t.Fatal("expected an error for a complex pencil")
	}
	bad, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6}, 2, 3)
	if _, _, err := EigenGeneralised(a, bad); err == nil {
		t.Fatal("expected an error for mismatched sizes")
	}
	rect, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6}, 3, 2)
	if _, _, err := EigenGeneralised(rect, b); err == nil {
		t.Fatal("expected an error for a rectangular a")
	}
	// b = all-ones is symmetric but singular: the Cholesky must refuse.
	singular, _ := core.FromFloats([]float64{1, 1, 1, 1}, 2, 2)
	if _, _, err := EigenGeneralised(mustFloats(t, []float64{1, 0, 0, 1}, 2, 2), singular); err == nil {
		t.Fatal("expected an error for a singular b")
	}
}