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

package linalg

import (
	"math"
	"testing"

	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

// mustF builds a float array or fails the test.
func mustF(t *testing.T, vals []float64, shape ...int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, shape...)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	return a
}

// TestPipelineHappyChain pins the values of a chained transformation
// against the same calls made directly.
func TestPipelineHappyChain(t *testing.T) {
	a := mustF(t, []float64{0.25, 4, 9, 16}, 2, 2)
	b := mustF(t, []float64{1, 1, 1, 1}, 2, 2)
	got, err := Pipe(a).AddF(1).Sqrt().Mul(b).Transpose().Reshape(4).Result()
	if err != nil {
		t.Fatalf("pipeline: %v", err)
	}
	want := core.Transpose(mustF(t, []float64{math.Sqrt(1.25), math.Sqrt(5), math.Sqrt(10), math.Sqrt(17)}, 2, 2))
	if got.Len() != 4 {
		t.Fatalf("result length %d, want 4", got.Len())
	}
	for i := range 4 {
		if math.Abs(got.FloatAt(i)-want.FloatAt(i)) > 1e-12 {
			t.Fatalf("[%d] = %g, want %g", i, got.FloatAt(i), want.FloatAt(i))
		}
	}
}

// TestPipelineShortCircuit pins the error contract: a failing step
// records the first error, Result returns the array that triggered it,
// and every later step is a no-op however wrong its arguments are.
func TestPipelineShortCircuit(t *testing.T) {
	a := mustF(t, []float64{1, 2, 3, 4}, 2, 2)
	// Reshape to an incompatible count fails.
	out, err := Pipe(a).Reshape(3, 5).Sqrt().AddF(1).MatMul2D(nil).Result()
	if err == nil {
		t.Fatal("expected the reshape error to surface")
	}
	if out != a {
		t.Fatal("Result must return the array that triggered the error")
	}
	// A singular Inv fails and freezes the chain.
	singular := mustF(t, []float64{1, 2, 2, 4}, 2, 2)
	out2, err2 := Pipe(singular).Inv().MulF(3).Result()
	if err2 == nil {
		t.Fatal("expected the singular Inv error to surface")
	}
	if out2 != singular {
		t.Fatal("Result must return the singular matrix itself")
	}
	// Reducing the only dimension of a 1-D array fails and freezes the
	// chain.
	vec := mustF(t, []float64{1, 2}, 2)
	out3, err3 := Pipe(vec).SumAxis(0).Exp().Result()
	if err3 == nil {
		t.Fatal("expected the SumAxis error to surface")
	}
	if out3 != vec {
		t.Fatal("Result must return the array the SumAxis failed on")
	}
}

// TestPipelineScalarsAndExtrema pins the scalar and reduction steps
// against the same calls made directly.
func TestPipelineScalarsAndExtrema(t *testing.T) {
	a := mustF(t, []float64{1, -2, 3, -4}, 2, 2)
	got, err := Pipe(a).MulF(2).ClipF(-5, 5).SumAxis(0).Result()
	if err != nil {
		t.Fatalf("pipeline: %v", err)
	}
	mul := core.MulF(a, 2)
	clipped, err := core.ClipF(mul, -5, 5)
	if err != nil {
		t.Fatalf("ClipF: %v", err)
	}
	want, err := core.SumAxis(clipped, 0)
	if err != nil {
		t.Fatalf("SumAxis: %v", err)
	}
	if got.Len() != want.Len() {
		t.Fatalf("length %d, want %d", got.Len(), want.Len())
	}
	for i := range want.Len() {
		if got.FloatAt(i) != want.FloatAt(i) {
			t.Fatalf("[%d] = %g, want %g", i, got.FloatAt(i), want.FloatAt(i))
		}
	}
	b := mustF(t, []float64{5, -6, 7, -8}, 4)
	got2, err2 := Pipe(b).Abs().Neg().Result()
	if err2 != nil {
		t.Fatalf("pipeline: %v", err2)
	}
	for i := range 4 {
		if got2.FloatAt(i) != float64(-i-5) {
			t.Fatalf("[%d] = %g", i, got2.FloatAt(i))
		}
	}
}