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

package optim

import (
	"math"
	"testing"

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

// Benchmarks for the fit surface of LevenbergMarquardtFit: the weighted
// residual path and the covariance report, the two options the plain
// entry point and its benchmarks never touch.

// benchExponentialFit builds the two-parameter decay fit the older LM
// benchmarks run, with the observations and their variances handed
// back so the weighted and covariance variants measure the same model.
func benchExponentialFit() (residual func(*core.Array) (*core.Array, error), jacobian func(*core.Array) (*core.Array, error), p0 *core.Array) {
	const nObs = 40
	t := make([]float64, nObs)
	y := make([]float64, nObs)
	for i := range nObs {
		t[i] = float64(i) / 4
		y[i] = 2.5*math.Exp(-0.7*t[i]) + 0.02*math.Sin(float64(i))
	}
	residual = func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, nObs)
		vals := out.RawFloats()
		for i := range nObs {
			vals[i] = p.FloatAt(0)*math.Exp(-p.FloatAt(1)*t[i]) - y[i]
		}
		return out, nil
	}
	jacobian = func(p *core.Array) (*core.Array, error) {
		out := core.New(core.Float, nObs, 2)
		vals := out.RawFloats()
		for i := range nObs {
			e := math.Exp(-p.FloatAt(1) * t[i])
			vals[i*2] = e
			vals[i*2+1] = -p.FloatAt(0) * t[i] * e
		}
		return out, nil
	}
	p0, _ = core.FromFloats([]float64{1, 0.2}, 2)
	return residual, jacobian, p0
}

// BenchmarkLevenbergMarquardtFitWeighted measures the fit with one
// variance per residual: the whitening of the residual and of the
// finite-difference Jacobian rides every evaluation the run makes.
func BenchmarkLevenbergMarquardtFitWeighted(b *testing.B) {
	residual, _, p0 := benchExponentialFit()
	const nObs = 40
	variance := make([]float64, nObs)
	for i := range nObs {
		variance[i] = 1 + float64(i)/8
	}
	sigma, err := core.FromFloats(variance, nObs)
	if err != nil {
		b.Fatal(err)
	}
	opts := LMOptions{MaxIterations: 30, Sigma: sigma}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := LevenbergMarquardtFit(residual, p0, opts); err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkLevenbergMarquardtFitCovariance measures the covariance
// report: the fit converges and then rebuilds the Jacobian at the
// answer, forms the normal equations from it and solves the identity
// through them.
func BenchmarkLevenbergMarquardtFitCovariance(b *testing.B) {
	residual, jacobian, p0 := benchExponentialFit()
	opts := LMOptions{MaxIterations: 30, Jacobian: jacobian, RequestCovariance: true}
	b.ReportAllocs()
	for b.Loop() {
		if _, err := LevenbergMarquardtFit(residual, p0, opts); err != nil {
			b.Fatal(err)
		}
	}
}