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

package optim

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

// TestLevenbergMarquardt fits y = a·e^{−bx} + c to exact data and
// checks that the fitted parameters match the generating values
// within the convergence tolerance of the LM fitter.
func TestLevenbergMarquardt(t *testing.T) {
	xData := []float64{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
	yData := make([]float64, len(xData))
	for i, x := range xData {
		yData[i] = 3*math.Exp(-0.5*float64(x)) + 0.5
	}
	yArr := mustFloats(t, yData, len(yData))
	residual := func(p *core.Array) (*core.Array, error) {
		a, b, c := p.FloatAt(0), p.FloatAt(1), p.FloatAt(2)
		out := core.New(core.Float, len(xData))
		for i := range len(xData) {
			out.RawFloats()[i] = yArr.FloatAt(i) - (a*math.Exp(-b*float64(xData[i])) + c)
		}
		return out, nil
	}
	p0 := mustFloats(t, []float64{2, 0.3, 0.1}, 3)
	pOpt, chi2, err := LevenbergMarquardt(residual, p0, LMOptions{})
	if err != nil {
		t.Fatalf("LevenbergMarquardt: %v", err)
	}
	if chi2 > 1e-4 {
		t.Fatalf("χ² = %v, want < 1e-4 for exact data", chi2)
	}
	a := pOpt.FloatAt(0)
	b := pOpt.FloatAt(1)
	if math.Abs(a-3) > 0.01 || math.Abs(b-0.5) > 0.01 {
		t.Fatalf("a = %.8f, b = %.8f, want ≈ 3, 0.5", a, b)
	}
}

// TestLevenbergMarquardtErrors pins the error contract.
func TestLevenbergMarquardtErrors(t *testing.T) {
	f := func(p *core.Array) (*core.Array, error) { return nil, nil }
	cx, _ := core.FromComplexes([]complex128{1}, 1)
	if _, _, err := LevenbergMarquardt(f, cx, LMOptions{}); err == nil {
		t.Fatal("expected an error for complex parameters")
	}
	empty, _ := core.FromFloats(nil, 0)
	if _, _, err := LevenbergMarquardt(f, empty, LMOptions{}); err == nil {
		t.Fatal("expected an error for an empty parameter vector")
	}
}

// TestLevenbergMarquardtAnalyticJacobian fits the same exponential
// model twice, once with the analytic Jacobian and once with the
// central-difference fallback: both must land on the same optimum,
// the analytic one reaching it without the extra evaluations.
func TestLevenbergMarquardtAnalyticJacobian(t *testing.T) {
	xData := []float64{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}
	yData := make([]float64, len(xData))
	for i, x := range xData {
		yData[i] = 3*math.Exp(-0.5*float64(x)) + 0.5
	}
	residual := func(p *core.Array) (*core.Array, error) {
		a, b, c := p.FloatAt(0), p.FloatAt(1), p.FloatAt(2)
		out := core.New(core.Float, len(xData))
		for i := range len(xData) {
			out.RawFloats()[i] = yData[i] - (a*math.Exp(-b*float64(xData[i])) + c)
		}
		return out, nil
	}
	jacobian := func(p *core.Array) (*core.Array, error) {
		a, b := p.FloatAt(0), p.FloatAt(1)
		out := core.New(core.Float, len(xData), 3)
		for i := range len(xData) {
			e := math.Exp(-b * float64(xData[i]))
			out.RawFloats()[i*3+0] = -e
			out.RawFloats()[i*3+1] = a * float64(xData[i]) * e
			out.RawFloats()[i*3+2] = -1
		}
		return out, nil
	}
	p0 := mustFloats(t, []float64{2, 0.3, 0.1}, 3)
	pA, chi2A, err := LevenbergMarquardt(residual, p0, LMOptions{Jacobian: jacobian})
	if err != nil {
		t.Fatalf("LevenbergMarquardt analytic: %v", err)
	}
	pD, chi2D, err := LevenbergMarquardt(residual, p0, LMOptions{})
	if err != nil {
		t.Fatalf("LevenbergMarquardt differences: %v", err)
	}
	for j := range 3 {
		if math.Abs(pA.FloatAt(j)-pD.FloatAt(j)) > 1e-6 {
			t.Fatalf("parameter %d: analytic %.12g, differences %.12g", j, pA.FloatAt(j), pD.FloatAt(j))
		}
	}
	if chi2A > 1e-4 {
		t.Fatalf("analytic χ² = %v, want < 1e-4", chi2A)
	}
	if math.Abs(chi2A-chi2D) > 1e-8 {
		t.Fatalf("χ² disagree: analytic %v, differences %v", chi2A, chi2D)
	}
}

// TestLevenbergMarquardtAnalyticExact solves a determined linear
// system, where the Gauss-Newton step is exact and the analytic
// Jacobian reaches the solution the equations dictate.
func TestLevenbergMarquardtAnalyticExact(t *testing.T) {
	residual := func(p *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{
			p.FloatAt(0) + 2*p.FloatAt(1) - 3,
			2*p.FloatAt(0) + p.FloatAt(1) - 4,
		}, 2)
	}
	jacobian := func(p *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{1, 2, 2, 1}, 2, 2)
	}
	p0 := mustFloats(t, []float64{0, 0}, 2)
	pOpt, chi2, err := LevenbergMarquardt(residual, p0, LMOptions{Jacobian: jacobian, Tolerance: 1e-14})
	if err != nil {
		t.Fatalf("LevenbergMarquardt: %v", err)
	}
	if math.Abs(pOpt.FloatAt(0)-5.0/3) > 1e-8 || math.Abs(pOpt.FloatAt(1)-2.0/3) > 1e-8 {
		t.Fatalf("p = (%.12g, %.12g), want (5/3, 2/3)", pOpt.FloatAt(0), pOpt.FloatAt(1))
	}
	if chi2 > 1e-16 {
		t.Fatalf("χ² = %v, want 0", chi2)
	}
}

// TestLevenbergMarquardtJacobianErrors pins the analytic Jacobian's
// own contract: shape mismatches and callback errors surface.
func TestLevenbergMarquardtJacobianErrors(t *testing.T) {
	residual := func(p *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{p.FloatAt(0) - 1, p.FloatAt(1) - 2}, 2)
	}
	rank1 := func(p *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{1, 1}, 2)
	}
	if _, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2),
		LMOptions{Jacobian: rank1}); err == nil {
		t.Fatal("expected an error for a rank-1 Jacobian")
	}
	wrongDims := func(p *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{1, 1, 1}, 3, 1)
	}
	if _, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2),
		LMOptions{Jacobian: wrongDims}); err == nil {
		t.Fatal("expected an error for a Jacobian of the wrong dimensions")
	}
	boom := func(p *core.Array) (*core.Array, error) {
		return nil, base.Errf("jacobian failed")
	}
	if _, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2),
		LMOptions{Jacobian: boom}); err == nil {
		t.Fatal("expected the Jacobian error to propagate")
	}
}

// TestLevenbergMarquardtDtypes pins the parameter-vector promotion:
// Int and Float32 starts must behave exactly like Float64 ones. The
// old RawFloats() copy silently started those fits from zeros.
func TestLevenbergMarquardtDtypes(t *testing.T) {
	residual := func(p *core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{
			p.FloatAt(0) + 2*p.FloatAt(1) - 3,
			2*p.FloatAt(0) + p.FloatAt(1) - 4,
		}, 2)
	}
	jacobian := func(*core.Array) (*core.Array, error) {
		return core.FromFloats([]float64{1, 2, 2, 1}, 2, 2)
	}
	ints, err := core.FromInts([]int64{0, 0}, 2)
	if err != nil {
		t.Fatalf("FromInts: %v", err)
	}
	thirtyTwo, err := core.FromFloat32s([]float32{0, 0}, 2)
	if err != nil {
		t.Fatalf("FromFloat32s: %v", err)
	}
	for name, p0 := range map[string]*core.Array{"int": ints, "float32": thirtyTwo} {
		pOpt, chi2, err := LevenbergMarquardt(residual, p0, LMOptions{Jacobian: jacobian, Tolerance: 1e-14})
		if err != nil {
			t.Fatalf("LevenbergMarquardt(%s start): %v", name, err)
		}
		if math.Abs(pOpt.FloatAt(0)-5.0/3) > 1e-8 || math.Abs(pOpt.FloatAt(1)-2.0/3) > 1e-8 {
			t.Fatalf("LevenbergMarquardt(%s start): p = (%.12g, %.12g), want (5/3, 2/3)",
				name, pOpt.FloatAt(0), pOpt.FloatAt(1))
		}
		if chi2 > 1e-16 {
			t.Fatalf("LevenbergMarquardt(%s start): χ² = %v, want 0", name, chi2)
		}
	}
}

// TestLevenbergMarquardtResidualLengthChange pins the residual
// contract: a callback whose output length changes mid-fit is an
// error naming the mismatch, not a shape panic.
func TestLevenbergMarquardtResidualLengthChange(t *testing.T) {
	calls := 0
	residual := func(*core.Array) (*core.Array, error) {
		calls++
		if calls == 1 {
			return core.FromFloats([]float64{1, 2, 3}, 3)
		}
		return core.FromFloats([]float64{1, 2}, 2)
	}
	_, _, err := LevenbergMarquardt(residual, mustFloats(t, []float64{0, 0}, 2), LMOptions{})
	if err == nil {
		t.Fatal("residual length change mid-fit: want an error")
	}
}

// TestLevenbergMarquardtStencilMatchesAnalyticJacobian pins the reused
// difference stencil against the analytic route. The stencil carries
// the offset on one parameter at a time and puts it back as soon as the
// column is differenced, so each column's Jacobian is the central
// difference of that parameter alone and the two routes land on the
// same point. A stencil that leaves its offset behind differences every
// later column at a point that is also displaced earlier: a different
// Jacobian, and a fit that parts company with the analytic route.
func TestLevenbergMarquardtStencilMatchesAnalyticJacobian(t *testing.T) {
	// A decay with an offset and a smooth nuisance term, so every
	// parameter carries curvature and the columns are coupled.
	const n = 40
	xs := make([]float64, n)
	ys := make([]float64, n)
	for i := range n {
		xs[i] = float64(i) * 0.25
		ys[i] = 1.5*math.Exp(-0.8*xs[i]) + 0.7 + 0.05*math.Sin(2*xs[i])
	}
	residual := func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		v := out.RawFloats()
		for i := range n {
			v[i] = ys[i] - (p.FloatAt(0)*math.Exp(-p.FloatAt(1)*xs[i]) +
				p.FloatAt(2) + p.FloatAt(3)*math.Sin(2*xs[i]))
		}
		return out, nil
	}
	jacobian := func(p *core.Array) (*core.Array, error) {
		v := make([]float64, n*4)
		for i := range n {
			e := math.Exp(-p.FloatAt(1) * xs[i])
			v[i*4+0] = -e
			v[i*4+1] = p.FloatAt(0) * xs[i] * e
			v[i*4+2] = -1
			v[i*4+3] = -math.Sin(2 * xs[i])
		}
		return core.FromFloats(v, n, 4)
	}
	start, err := core.FromFloats([]float64{1.2, 0.9, 0.4, 0.04}, 4)
	if err != nil {
		t.Fatal(err)
	}
	fd, fdChi, err := LevenbergMarquardt(residual, start, LMOptions{})
	if err != nil {
		t.Fatalf("LevenbergMarquardt with a difference stencil: %v", err)
	}
	analytic, analyticChi, err := LevenbergMarquardt(residual, start, LMOptions{Jacobian: jacobian})
	if err != nil {
		t.Fatalf("LevenbergMarquardt with an analytic Jacobian: %v", err)
	}
	if fd.Len() != analytic.Len() {
		t.Fatalf("the two routes fit %d and %d parameters", fd.Len(), analytic.Len())
	}
	for i := range fd.Len() {
		if d := math.Abs(fd.FloatAt(i) - analytic.FloatAt(i)); d > 1e-12 {
			t.Fatalf("parameter %d: the stencil route gives %.17g, the analytic one %.17g (differ by %g)",
				i, fd.FloatAt(i), analytic.FloatAt(i), d)
		}
	}
	if math.Abs(fdChi-analyticChi) > 1e-12 {
		t.Fatalf("the two routes report chi2 %.17g and %.17g", fdChi, analyticChi)
	}
	// The stencil is carried across the fit's iterations, so the same
	// start must reproduce the same bits.
	again, againChi, err := LevenbergMarquardt(residual, start, LMOptions{})
	if err != nil {
		t.Fatalf("the repeated stencil fit: %v", err)
	}
	for i := range again.Len() {
		if math.Float64bits(again.FloatAt(i)) != math.Float64bits(fd.FloatAt(i)) {
			t.Fatalf("parameter %d moved between stencil fits: %.17g against %.17g", i, again.FloatAt(i), fd.FloatAt(i))
		}
	}
	if math.Float64bits(againChi) != math.Float64bits(fdChi) {
		t.Fatalf("chi2 moved between stencil fits: %.17g against %.17g", againChi, fdChi)
	}
}