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

package integrate

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

import (
	"math"
	"strings"
	"testing"
)

// daeCircuit returns the source, mass matrix and state of a linear
// index-1 circuit: a one-volt source feeds a unit resistor into node
// v1 (unit capacitor to ground), an inductor of one henry carries i
// on to node v2, and node v2 dumps through a unit resistor with no
// capacitor, so its KCL row 0 = i − v2 is the algebraic constraint
// and i the algebraic variable. With C = L = R = 1 the differential
// pair is x' = Ax + (1, 0) with A = [[−1, −1], [1, −1]], whose
// solution is elementary.
func daeCircuit(t float64, y *core.Array) (*core.Array, error) {
	return core.FromFloats([]float64{
		1 - y.FloatAt(0) - y.FloatAt(2),
		y.FloatAt(2) - y.FloatAt(1),
		y.FloatAt(0) - y.FloatAt(1),
	}, 3)
}

func daeCircuitEnd(t *testing.T, steps int, y0 []float64, t0, t1 float64) []float64 {
	t.Helper()
	m := mustFloats(t, []float64{1, 0, 0, 0, 0, 0, 0, 0, 1}, 3, 3)
	end, err := IntegrateDAE(daeCircuit, m, t0, t1, mustFloats(t, y0), steps, DAEOptions{})
	if err != nil {
		t.Fatalf("IntegrateDAE: %v", err)
	}
	return []float64{end.FloatAt(0), end.FloatAt(1), end.FloatAt(2)}
}

// daeCircuitExact evaluates the exact v1, v2, i at time t: the
// equilibrium (0.5, 0.5) plus the elementary homogeneous part.
func daeCircuitExact(t float64) []float64 {
	c := math.Exp(-t) / 2
	return []float64{
		0.5 + c*(math.Cos(t)+math.Sin(t)),
		0.5 + c*(math.Sin(t)-math.Cos(t)),
		0.5 + c*(math.Sin(t)-math.Cos(t)),
	}
}

// TestIntegrateDAECircuit is the linear index-1 pin: the differential
// nodes track the elementary solution, and the algebraic variable i
// satisfies the KCL constraint i = v2 to rounding at the end state,
// because every step enforces the constraint row exactly.
func TestIntegrateDAECircuit(t *testing.T) {
	end := daeCircuitEnd(t, 200, []float64{1, 0, 0}, 0, 1)
	want := daeCircuitExact(1)
	for k, band := range []float64{0.01, 0.01, 0.01} {
		if math.Abs(end[k]-want[k]) > band {
			t.Fatalf("circuit[%d] = %.14g, want %.14g ± %g", k, end[k], want[k], band)
		}
	}
	if math.Abs(end[1]-end[2]) > 1e-10 {
		t.Fatalf("the constraint i = v2 drifted to %g at the end state", end[1]-end[2])
	}
}

// TestIntegrateDAEScalarConstraint pins the algebraic variable on the
// exact constraint to rounding: with w' unconstrained by M's zero row
// and 0 = w − cos t, the solved w must equal cos at every step, so
// certainly at the end.
func TestIntegrateDAEScalarConstraint(t *testing.T) {
	f := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-y.FloatAt(0), y.FloatAt(1) - math.Cos(t)}, 2)
	}
	m := mustFloats(t, []float64{1, 0, 0, 0}, 2, 2)
	end, err := IntegrateDAE(f, m, 0, 1, mustFloats(t, []float64{1, 1}), 25, DAEOptions{})
	if err != nil {
		t.Fatalf("IntegrateDAE: %v", err)
	}
	if math.Abs(end.FloatAt(1)-math.Cos(1)) > 1e-12 {
		t.Fatalf("algebraic w(1) = %.16g, want cos(1) = %.16g to rounding",
			end.FloatAt(1), math.Cos(1))
	}
	if math.Abs(end.FloatAt(0)-math.Exp(-1)) > 0.05 {
		t.Fatalf("differential u(1) = %.14g, want %.14g ± 0.05", end.FloatAt(0), math.Exp(-1))
	}
}

// TestIntegrateDAEBackward integrates the circuit backwards from the
// exact end state; the signed-step formulation must return the start.
func TestIntegrateDAEBackward(t *testing.T) {
	want := daeCircuitExact(1)
	end := daeCircuitEnd(t, 200, want, 1, 0)
	start := daeCircuitExact(0)
	for k := range 3 {
		if math.Abs(end[k]-start[k]) > 0.01 {
			t.Fatalf("backward circuit[%d] = %.14g, want %.14g ± 0.01", k, end[k], start[k])
		}
	}
}

// TestIntegrateDAEConsistencyRefused pins the initial-residual check:
// a start violating the KCL row by one full unit is refused, with the
// row named.
func TestIntegrateDAEConsistencyRefused(t *testing.T) {
	m := mustFloats(t, []float64{1, 0, 0, 0, 0, 0, 0, 0, 1}, 3, 3)
	_, err := IntegrateDAE(daeCircuit, m, 0, 1, mustFloats(t, []float64{1, 1, 0}), 200, DAEOptions{})
	if err == nil {
		t.Fatal("expected an error for an inconsistent start")
	}
	if !strings.Contains(err.Error(), "row 1") {
		t.Fatalf("want the algebraic row named, got %v", err)
	}
}

// TestIntegrateDAEPendulumRefused pins the honest index-3 refusal: the
// Cartesian pendulum with multipliers has a mass matrix that admits
// index 1 by rank alone, but its algebraic rows (the constraints) do
// not depend on the algebraic variables (the multipliers) at all, so
// the certified block is singular and the solver refuses, naming the
// detection.
func TestIntegrateDAEPendulumRefused(t *testing.T) {
	// y = (x, ypos, u, v, lambda, mu): position, velocity, multipliers.
	// m = 1, g = 0, length 1: the unit circle, started at (1, 0) with
	// unit tangential speed and the multiplier that holds it there.
	f := func(t float64, y *core.Array) (*core.Array, error) {
		x, ypos, u, v, lambda := y.FloatAt(0), y.FloatAt(1), y.FloatAt(2), y.FloatAt(3), y.FloatAt(4)
		return core.FromFloats([]float64{
			u, v,
			-2 * x * lambda,
			-2 * ypos * lambda,
			x*x + ypos*ypos - 1,
			x*u + ypos*v,
		}, 6)
	}
	m := mustFloats(t, []float64{
		1, 0, 0, 0, 0, 0,
		0, 1, 0, 0, 0, 0,
		0, 0, 1, 0, 0, 0,
		0, 0, 0, 1, 0, 0,
		0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0,
	}, 6, 6)
	y0 := mustFloats(t, []float64{1, 0, 0, 1, 0.5, 0})
	_, err := IntegrateDAE(f, m, 0, 0.1, y0, 10, DAEOptions{})
	if err == nil {
		t.Fatal("expected the index-3 pendulum to be refused")
	}
	if !strings.Contains(err.Error(), "index 1") || !strings.Contains(err.Error(), "singular") {
		t.Fatalf("want the index detection stated, got %v", err)
	}
}

// TestIntegrateDAEIndexTwoStall pins the other honest refusal: an
// ordinary stiff ODE whose per-step Newton matrix is exactly singular
// at the chosen step (y0' = 100·y0 with h·100 = 1) fails loudly
// through the Newton solve, not through silent drift; the index-1
// certificate itself passes because the algebraic row y1 − y0 does
// depend on the algebraic variable, so the refusal here comes from
// the differential row's pathology and must be named as such.
func TestIntegrateDAEIndexTwoStall(t *testing.T) {
	// y0' = 100·y0 with h·100 = 1 makes the Newton matrix singular.
	f := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{100 * y.FloatAt(0), y.FloatAt(1) - y.FloatAt(0)}, 2)
	}
	m := mustFloats(t, []float64{1, 0, 0, 0}, 2, 2)
	_, err := IntegrateDAE(f, m, 0, 1, mustFloats(t, []float64{1, 1}), 100, DAEOptions{})
	if err == nil {
		t.Fatal("expected the singular per-step solve to be refused")
	}
	if !strings.Contains(err.Error(), "Newton") {
		t.Fatalf("want a Newton failure, got %v", err)
	}
}

// TestIntegrateDAEErrors pins the structural error contract: a zero
// step count, a rank-1 mass matrix of the wrong shape, a nonsingular
// matrix, a rank deficiency without whole zero rows, mismatched zero
// row and column counts, an empty state, a non-finite matrix entry and
// a failing f are all errors; a degenerate span returns the start.
func TestIntegrateDAEErrors(t *testing.T) {
	good := mustFloats(t, []float64{1, 0, 0, 0}, 2, 2)
	simple := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-y.FloatAt(0), y.FloatAt(1)}, 2)
	}
	simple3 := func(t float64, y *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{-y.FloatAt(0), y.FloatAt(1), -y.FloatAt(2)}, 3)
	}
	if _, err := IntegrateDAE(simple, good, 0, 1, mustFloats(t, []float64{1, 1}), 0, DAEOptions{}); err == nil {
		t.Fatal("expected an error for zero steps")
	}
	badShape := mustFloats(t, []float64{1, 0}, 1, 2)
	if _, err := IntegrateDAE(simple, badShape, 0, 1, mustFloats(t, []float64{1, 1}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for a non-square mass matrix")
	}
	identity := mustFloats(t, []float64{1, 0, 0, 1}, 2, 2)
	if _, err := IntegrateDAE(simple, identity, 0, 1, mustFloats(t, []float64{1, 1}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for a nonsingular mass matrix")
	}
	noZeroRows := mustFloats(t, []float64{1, 1, 1, 1}, 2, 2)
	if _, err := IntegrateDAE(simple, noZeroRows, 0, 1, mustFloats(t, []float64{1, 1}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for rank deficiency without zero rows")
	}
	rowColMismatch := mustFloats(t, []float64{1, 1, 0, 0}, 2, 2)
	if _, err := IntegrateDAE(simple, rowColMismatch, 0, 1, mustFloats(t, []float64{1, 1}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for mismatched zero row and column counts")
	}
	if _, err := IntegrateDAE(simple, good, 0, 1, mustFloats(t, nil), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for an empty state")
	}
	nonFinite := mustFloats(t, []float64{1, 0, 0, math.Inf(1)}, 2, 2)
	if _, err := IntegrateDAE(simple, nonFinite, 0, 1, mustFloats(t, []float64{1, 1}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for a non-finite mass matrix entry")
	}
	complexM, _ := core.FromComplexes([]complex128{1, 0, 0, 1}, 2, 2)
	if _, err := IntegrateDAE(simple, complexM, 0, 1, mustFloats(t, []float64{1, 1}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for a complex mass matrix")
	}
	// One zero row but a second dependent row: the rank deficiency
	// exceeds the zero rows and the contract is refused.
	hiddenDeficiency := mustFloats(t, []float64{1, 1, 0, 1, 1, 0, 0, 0, 0}, 3, 3)
	_, err := IntegrateDAE(simple3, hiddenDeficiency, 0, 1, mustFloats(t, []float64{1, 1, 1}), 10, DAEOptions{})
	if err == nil || !strings.Contains(err.Error(), "rank deficiency") {
		t.Fatalf("expected the hidden rank deficiency to be refused, got %v", err)
	}
	// A degenerate span answers the validated start unchanged. The
	// start must satisfy the algebraic row of this system, y1 = 0.
	same, err := IntegrateDAE(simple, good, 1, 1, mustFloats(t, []float64{1, 0}), 10, DAEOptions{})
	if err != nil {
		t.Fatalf("zero span: %v", err)
	}
	if same.FloatAt(0) != 1 || same.FloatAt(1) != 0 {
		t.Fatalf("zero span moved the state to (%v, %v)", same.FloatAt(0), same.FloatAt(1))
	}
	boom := func(t float64, y *core.Array) (*core.Array, error) {
		if t > 0.5 {
			return nil, base.Errf("detector tripped")
		}
		return core.FromFloats([]float64{-y.FloatAt(0), y.FloatAt(1)}, 2)
	}
	if _, err := IntegrateDAE(boom, good, 0, 1, mustFloats(t, []float64{1, 0}), 100, DAEOptions{}); err == nil {
		t.Fatal("expected the operator error to propagate")
	}
	// An f failing on the initial evaluation, the initial Jacobian and
	// inside the first Newton iteration is refused at once.
	always := func(t float64, y *core.Array) (*core.Array, error) {
		return nil, base.Errf("detector tripped")
	}
	if _, err := IntegrateDAE(always, good, 0, 1, mustFloats(t, []float64{1, 0}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected an error for an f that always fails")
	}
	// An f that only tolerates the exact seed fails when the Newton
	// iteration perturbs the state for its numerical Jacobian.
	touchy := func(t float64, y *core.Array) (*core.Array, error) {
		if y.FloatAt(0) != 1 {
			return nil, base.Errf("detector tripped")
		}
		return core.FromFloats([]float64{1 - y.FloatAt(0), y.FloatAt(1) - y.FloatAt(0)}, 2)
	}
	if _, err := IntegrateDAE(touchy, good, 0, 1, mustFloats(t, []float64{1, 0}), 10, DAEOptions{}); err == nil {
		t.Fatal("expected the Jacobian perturbation to trip the f error")
	}
}