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

package optim

import (
	"fmt"
	"math"
	"runtime"
	"sync/atomic"
	"testing"
	"time"

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

// Tests and benchmarks for the opt-in parallel sweep of the
// central-difference Jacobian, LMOptions.ParallelJacobian and
// RootSystemOptions.ParallelJacobian. The contract under test: the
// default false never calls the callback from more than one goroutine
// and answers bit for bit what it answered before; the true value may
// overlap the callback's evaluations across columns and still answers
// bit for bit the same numbers, because every column is differenced by
// the same stencil against the same point.

// concurrencyTracker records the largest number of callbacks it has
// seen inside at once. The counter and the peak are atomic: the tracker
// itself must not be the thing that serialises the calls it measures.
type concurrencyTracker struct {
	cur  atomic.Int64
	peak atomic.Int64
}

func (c *concurrencyTracker) enter() {
	n := c.cur.Add(1)
	for {
		p := c.peak.Load()
		if n <= p || c.peak.CompareAndSwap(p, n) {
			return
		}
	}
}

func (c *concurrencyTracker) leave() { c.cur.Add(-1) }

// lmParallelProblem builds an eight-parameter two-exponential plus
// sinusoid fit with a known generating model, hard enough that the fit
// walks a number of damped steps before it converges.
func lmParallelProblem() (residual func(*core.Array) (*core.Array, error), start []float64) {
	const nObs = 48
	t := make([]float64, nObs)
	y := make([]float64, nObs)
	truth := []float64{2.0, 0.7, 1.5, 1.9, 0.8, 3.1, 0.4, 0.05}
	for i := range nObs {
		t[i] = float64(i) / 6
		y[i] = truth[0]*math.Exp(-truth[1]*t[i]) + truth[2]*math.Exp(-truth[3]*t[i]) +
			truth[4]*math.Sin(truth[5]*t[i]+truth[6]) + truth[7]*t[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]) + p.FloatAt(2)*math.Exp(-p.FloatAt(3)*t[i]) +
				p.FloatAt(4)*math.Sin(p.FloatAt(5)*t[i]+p.FloatAt(6)) + p.FloatAt(7)*t[i] - y[i]
		}
		return out, nil
	}
	return residual, []float64{1.6, 0.5, 1.2, 1.6, 0.6, 2.7, 0.2, 0.02}
}

// rootParallelProblem builds eight coupled quadratic equations whose
// root sits near xᵢ = i+1, hard enough that Newton walks several rounds.
func rootParallelProblem() (f func(*core.Array) (*core.Array, error), start []float64) {
	const n = 8
	f = func(x *core.Array) (*core.Array, error) {
		out := core.New(core.Float, n)
		vals := out.RawFloats()
		for i := range n {
			vals[i] = x.FloatAt(i)*x.FloatAt(i) - float64(i+1)*float64(i+1) + 0.1*x.FloatAt((i+1)%n)
		}
		return out, nil
	}
	start = make([]float64, n)
	for i := range start {
		start[i] = float64(i+1) + 0.5
	}
	return f, start
}

func TestParallelJacobianBitIdenticalFit(t *testing.T) {
	residual, start := lmParallelProblem()
	run := func(parallel bool) (*core.Array, float64) {
		t.Helper()
		p, chi2, err := LevenbergMarquardt(residual, mustFloats(t, start), LMOptions{
			MaxIterations:    100,
			ParallelJacobian: parallel,
		})
		if err != nil {
			t.Fatalf("LevenbergMarquardt(ParallelJacobian=%v): %v", parallel, err)
		}
		return p, chi2
	}
	pSerial, chi2Serial := run(false)
	pParallel, chi2Parallel := run(true)
	if math.Float64bits(chi2Serial) != math.Float64bits(chi2Parallel) {
		t.Fatalf("chi2 differs: serial %v, parallel %v", chi2Serial, chi2Parallel)
	}
	serial, parallel := pSerial.RawFloats(), pParallel.RawFloats()
	for i := range serial {
		if math.Float64bits(serial[i]) != math.Float64bits(parallel[i]) {
			t.Fatalf("parameter %d differs: serial %.17g, parallel %.17g", i, serial[i], parallel[i])
		}
	}
}

func TestParallelJacobianBitIdenticalRoot(t *testing.T) {
	f, start := rootParallelProblem()
	run := func(parallel, broyden bool) (*core.Array, float64) {
		t.Helper()
		x, res, err := FindRootSystem(f, mustFloats(t, start), RootSystemOptions{
			MaxIterations:    100,
			UseBroyden:       broyden,
			ParallelJacobian: parallel,
		})
		if err != nil {
			t.Fatalf("FindRootSystem(UseBroyden=%v, ParallelJacobian=%v): %v", broyden, parallel, err)
		}
		return x, res
	}
	for _, broyden := range []bool{false, true} {
		xSerial, resSerial := run(false, broyden)
		xParallel, resParallel := run(true, broyden)
		if math.Float64bits(resSerial) != math.Float64bits(resParallel) {
			t.Fatalf("UseBroyden=%v: residual norm differs: serial %v, parallel %v", broyden, resSerial, resParallel)
		}
		sSerial, sParallel := xSerial.RawFloats(), xParallel.RawFloats()
		for i := range sSerial {
			if math.Float64bits(sSerial[i]) != math.Float64bits(sParallel[i]) {
				t.Fatalf("UseBroyden=%v: root coordinate %d differs: serial %.17g, parallel %.17g",
					broyden, i, sSerial[i], sParallel[i])
			}
		}
	}
}

// TestParallelJacobianBitIdenticalBroydenRebuilds drives the rebuild
// path: the duplicated equation x² = 1 has a rank-one Jacobian, so a
// Broyden run from the hopeless start rebuilds the numerical Jacobian
// every round until the budget refuses it. Serial and parallel must
// refuse with the same words, which pins every rebuilt Jacobian's bits.
func TestParallelJacobianBitIdenticalBroydenRebuilds(t *testing.T) {
	rankOne := func(x *core.Array) (*core.Array, error) {
		v := x.FloatAt(0)*x.FloatAt(0) - 1
		out := core.New(core.Float, 2)
		vals := out.RawFloats()
		vals[0], vals[1] = v, v
		return out, nil
	}
	run := func(parallel bool) string {
		t.Helper()
		_, _, err := FindRootSystem(rankOne, mustFloats(t, []float64{1.5, 1.5}), RootSystemOptions{
			MaxIterations:    12,
			UseBroyden:       true,
			ParallelJacobian: parallel,
		})
		if err == nil {
			t.Fatalf("FindRootSystem(ParallelJacobian=%v): want the budget refusal", parallel)
		}
		return err.Error()
	}
	if serial, parallel := run(false), run(true); serial != parallel {
		t.Fatalf("the refusal differs: serial %q, parallel %q", serial, parallel)
	}
}

// TestParallelJacobianCallbackConcurrency pins the consent boundary:
// with the option off the residual callback never runs inside more
// than one goroutine at once, and with it on the sweeps do overlap.
func TestParallelJacobianCallbackConcurrency(t *testing.T) {
	prev := engine.SetNumWorkers(4)
	defer engine.SetNumWorkers(prev)

	const nP = 8
	const nObs = 8
	peak := func(t *testing.T, parallel bool) int64 {
		t.Helper()
		var tr concurrencyTracker
		residual := func(p *core.Array) (*core.Array, error) {
			tr.enter()
			defer tr.leave()
			if parallel {
				// The caller holds its ground until a second one is
				// inside, so the overlap the option promises is
				// demonstrated by construction instead of left to the
				// scheduler to interleave two calls this short. A
				// sweep that collapsed to one goroutine would wait
				// out the grace alone and fail the peak below.
				deadline := time.Now().Add(250 * time.Millisecond)
				for tr.peak.Load() < 2 && time.Now().Before(deadline) {
					runtime.Gosched()
				}
			}
			s := 0.0
			for k := range 800 {
				s += math.Sin(float64(k)*7e-4+p.FloatAt(k&(nP-1))) * math.Cos(float64(k)*3e-4)
			}
			out := core.New(core.Float, nObs)
			vals := out.RawFloats()
			for i := range nObs {
				vals[i] = s + p.FloatAt(i&(nP-1))
			}
			return out, nil
		}
		start := make([]float64, nP)
		if _, _, err := LevenbergMarquardt(residual, mustFloats(t, start), LMOptions{
			MaxIterations:    2,
			AllowBudgetExit:  true,
			ParallelJacobian: parallel,
		}); err != nil {
			t.Fatalf("LevenbergMarquardt(ParallelJacobian=%v): %v", parallel, err)
		}
		return tr.peak.Load()
	}
	t.Run("serial", func(t *testing.T) {
		if got := peak(t, false); got != 1 {
			t.Fatalf("the serial run reached %d concurrent callback calls, want exactly 1", got)
		}
	})
	t.Run("parallel", func(t *testing.T) {
		if got := peak(t, true); got < 2 {
			t.Fatalf("the parallel run reached %d concurrent callback calls, want more than 1", got)
		}
	})
}

// TestParallelJacobianErrorMatchesSerial fails one column's stencil
// and requires the parallel sweep to report the same error the serial
// walk reports: the lowest failing column, not whichever worker got
// there first.
func TestParallelJacobianErrorMatchesSerial(t *testing.T) {
	const nP = 6
	const nObs = 10
	// Coordinate 3 is nonzero exactly while column 3's stencils run,
	// so only that column's residual evaluations fail.
	residual := func(p *core.Array) (*core.Array, error) {
		if p.FloatAt(3) != 0 {
			return nil, fmt.Errorf("the stencil touched column 3")
		}
		out := core.New(core.Float, nObs)
		for i := range nObs {
			out.RawFloats()[i] = p.FloatAt(0) + float64(i)
		}
		return out, nil
	}
	run := func(parallel bool) string {
		t.Helper()
		start := make([]float64, nP)
		_, _, err := LevenbergMarquardt(residual, mustFloats(t, start), LMOptions{
			MaxIterations:    10,
			ParallelJacobian: parallel,
		})
		if err == nil {
			t.Fatalf("LevenbergMarquardt(ParallelJacobian=%v): want the stencil error", parallel)
		}
		return err.Error()
	}
	serial := run(false)
	parallel := run(true)
	if serial != parallel {
		t.Fatalf("the error differs: serial %q, parallel %q", serial, parallel)
	}
}

// benchJacResidual builds a sixteen-parameter least squares problem
// whose residual costs roughly twenty microseconds of real arithmetic
// per call, and whose data sit far from every model the fit can reach,
// so a run spends a fixed number of rounds, almost all of it inside
// the Jacobian sweep, and never converges.
func benchJacResidual(b *testing.B) (residual func(*core.Array) (*core.Array, error), start *core.Array) {
	b.Helper()
	const nP = 16
	const nObs = 64
	const spin = 512
	t := make([]float64, nObs)
	y := make([]float64, nObs)
	for i := range nObs {
		t[i] = float64(i) / 8
		y[i] = 1e3 + float64(i)
	}
	residual = func(p *core.Array) (*core.Array, error) {
		s := 0.0
		for k := range spin {
			s += math.Sin(float64(k)*7e-4+p.FloatAt(k&(nP-1))) * math.Cos(float64(k)*3e-4)
		}
		out := core.New(core.Float, nObs)
		vals := out.RawFloats()
		for i := range nObs {
			v := s
			for k := range nP / 2 {
				v += p.FloatAt(2*k) * math.Exp(-p.FloatAt(2*k+1)*t[i])
			}
			vals[i] = v - y[i]
		}
		return out, nil
	}
	p0 := make([]float64, nP)
	for i := range p0 {
		p0[i] = 0.5
	}
	arr, err := core.FromFloats(p0, nP)
	if err != nil {
		b.Fatal(err)
	}
	return residual, arr
}

// BenchmarkLevenbergMarquardtJacobianSweep times the central-difference
// Jacobian sweep serial against parallel in one binary, the same fit
// walked both ways. The worker count mirrors the GOMAXPROCS=8 the run
// is read under: eight workers, one per thread.
func BenchmarkLevenbergMarquardtJacobianSweep(b *testing.B) {
	restore := engine.SetNumWorkers(8)
	defer engine.SetNumWorkers(restore)
	residual, start := benchJacResidual(b)
	for _, parallel := range []bool{false, true} {
		name := "serial"
		if parallel {
			name = "parallel"
		}
		b.Run(name, func(b *testing.B) {
			opts := LMOptions{MaxIterations: 12, AllowBudgetExit: true, ParallelJacobian: parallel}
			b.ReportAllocs()
			for b.Loop() {
				if _, _, err := LevenbergMarquardt(residual, start, opts); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}