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

// spdSample builds a deterministic symmetric positive definite matrix:
// B + Bᵀ + n·I from a fixed entry pattern, so no test randomness leaks.
func spdSample(n int) *core.Array {
	raw := make([]float64, n*n)
	for i := range n {
		for j := range n {
			raw[i*n+j] = math.Sin(float64(3*i+j+1)) + math.Cos(float64(i-j))
		}
	}
	vals := make([]float64, n*n)
	for i := range n {
		for j := range n {
			vals[i*n+j] = raw[i*n+j] + raw[j*n+i]
		}
		vals[i*n+i] += float64(n) // shift onto the PD cone
	}
	a, _ := core.FromFloats(vals, n, n)
	return a
}

// gramReconstruct multiplies a lower triangle out: returns L·Lᵀ.
func gramReconstruct(t *testing.T, l *core.Array) *core.Array {
	t.Helper()
	n := l.Shape()[0]
	out, err := zeros(core.Float, []int{n, n})
	if err != nil {
		t.Fatalf("gramReconstruct: %v", err)
	}
	for i := range n {
		for j := range i + 1 {
			s := 0.0
			for k := range n {
				s += l.FloatAt(i*n+k) * l.FloatAt(j*n+k)
			}
			out.SetFloatAt(i*n+j, s)
			out.SetFloatAt(j*n+i, s)
		}
	}
	return out
}

func matDiff(a, b *core.Array) float64 {
	worst := 0.0
	for i := range a.Len() {
		d := math.Abs(a.FloatAt(i) - b.FloatAt(i))
		if d > worst {
			worst = d
		}
	}
	return worst
}

// TestCholeskyUpdateGram pins the defining property: the rebuilt
// factor's Gram matrix equals A + x·xᵀ, checked against a fresh
// factorisation of the modified matrix.
func TestCholeskyUpdateGram(t *testing.T) {
	n := 5
	a := spdSample(n)
	l, err := Cholesky(a)
	if err != nil {
		t.Fatalf("Cholesky: %v", err)
	}
	x := mustFloats(t, []float64{1, -2, 0.5, 3, -1}, n)
	updated, err := CholeskyUpdate(l, x)
	if err != nil {
		t.Fatalf("CholeskyUpdate: %v", err)
	}
	want, err := zeros(core.Float, []int{n, n})
	if err != nil {
		t.Fatalf("zeros: %v", err)
	}
	for i := range n {
		for j := range n {
			want.SetFloatAt(i*n+j, a.FloatAt(i*n+j)+x.FloatAt(i)*x.FloatAt(j))
		}
	}
	got := gramReconstruct(t, updated)
	if matDiff(got, want) > 1e-10 {
		t.Fatalf("Gram mismatch %.3g", matDiff(got, want))
	}
	reference, err := Cholesky(want)
	if err != nil {
		t.Fatalf("Cholesky of the update: %v", err)
	}
	if matDiff(updated, reference) > 1e-8 {
		t.Fatalf("factors disagree %.3g", matDiff(updated, reference))
	}
}

// TestCholeskyUpdateDowndateRoundTrip updates and downdates the same
// vector: the original factor must come back.
func TestCholeskyUpdateDowndateRoundTrip(t *testing.T) {
	n := 6
	a := spdSample(n)
	l, err := Cholesky(a)
	if err != nil {
		t.Fatalf("Cholesky: %v", err)
	}
	x := mustFloats(t, []float64{0.3, 1, -0.7, 2, -1, 0.2}, n)
	up, err := CholeskyUpdate(l, x)
	if err != nil {
		t.Fatalf("CholeskyUpdate: %v", err)
	}
	back, err := CholeskyDowndate(up, x)
	if err != nil {
		t.Fatalf("CholeskyDowndate: %v", err)
	}
	if matDiff(back, l) > 1e-8 {
		t.Fatalf("round trip lost %.3g", matDiff(back, l))
	}
}

// TestCholeskyDowndateOutsideCone pins the honest failure: removing
// more than the matrix carries leaves the positive definite cone and
// the downdate must refuse.
func TestCholeskyDowndateOutsideCone(t *testing.T) {
	a := spdSample(3)
	l, err := Cholesky(a)
	if err != nil {
		t.Fatalf("Cholesky: %v", err)
	}
	big := mustFloats(t, []float64{10, 10, 10}, 3)
	if _, err := CholeskyDowndate(l, big); err == nil {
		t.Fatal("expected an error for a downdate outside the cone")
	}
	// A vector the matrix can absorb must succeed.
	small := mustFloats(t, []float64{0.1, 0.1, 0.1}, 3)
	if _, err := CholeskyDowndate(l, small); err != nil {
		t.Fatalf("small downdate: %v", err)
	}
}

// TestCholeskyRankOneErrors pins shape and dtype validation.
func TestCholeskyRankOneErrors(t *testing.T) {
	l, _ := Cholesky(spdSample(3))
	bad := mustFloats(t, []float64{1, 2}, 2)
	if _, err := CholeskyUpdate(l, bad); err == nil {
		t.Fatal("expected an error for a mismatched vector")
	}
	if _, err := CholeskyUpdate(bad, bad); err == nil {
		t.Fatal("expected an error for a non-square factor")
	}
	cx, _ := core.FromComplexes([]complex128{1}, 1)
	if _, err := CholeskyUpdate(l, cx); err == nil {
		t.Fatal("expected an error for a complex vector")
	}
}