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

package core

import (
	"math"

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

// Numerical Jacobians for vector-valued functions, the standalone
// companion of the central differences LevenbergMarquardt builds
// internally: one column per input coordinate, two evaluations per
// column, and a per-column step that scales with the coordinate's
// magnitude so every column carries the same relative resolution.

// JacobianOptions tunes Jacobian. Step is the absolute difference
// step applied to every coordinate; zero or negative selects the
// default per-column step sqrt(ε)·max(1, |x_j|), the largest step
// whose central-difference truncation error still sits below the
// rounding floor.
type JacobianOptions struct {
	Step float64
}

// Jacobian returns the Jacobian of f at x as an (m × n) float array,
// entry (i, j) holding ∂f_i/∂x_j by central differences with the
// step opts.Step, or sqrt(ε)·max(1, |x_j|) per column when unset. x
// may hold any real dtype and shape; its n elements are perturbed one
// at a time. f must map the point and every probe to a real rank-1
// array of one fixed length m: a non-vector output, an output that
// changes length between columns, an empty output, or an error from
// f is reported. The probe arrays handed to f are reused between
// columns, so f must not retain them.
func Jacobian(f func(x *Array) (*Array, error), x *Array, opts JacobianOptions) (*Array, error) {
	if x.dt == Complex {
		return nil, errf("Jacobian: complex points are not supported")
	}
	n := x.Len()
	if n == 0 {
		return nil, errf("Jacobian: the point must not be empty")
	}
	// The unperturbed evaluation fixes m up front and holds f to it
	// for every column that follows.
	base0, err := f(x)
	if err != nil {
		return nil, base.WrapErr("Jacobian", err)
	}
	if base0.dt == Complex {
		return nil, errf("Jacobian: complex outputs are not supported")
	}
	if base0.NDim() != 1 {
		return nil, errf("Jacobian: f must return a vector, got shape %s", shapeText(base0.Shape()))
	}
	m := base0.Len()
	if m == 0 {
		return nil, errf("Jacobian: f returns an empty vector")
	}
	out := &Array{shape: []int{m, n}, dt: Float}
	out.alloc(m * n)
	// The point widened once, then two probe copies reused across the
	// coordinate sweep: f receives an array it may keep for the
	// duration of the call, but each column rewrites both from
	// scratch.
	p := make([]float64, n)
	for k := range n {
		p[k] = x.floatAt(k)
	}
	pp := make([]float64, n)
	pm := make([]float64, n)
	for j := range n {
		h := opts.Step
		if h <= 0 {
			h = math.Sqrt(base.EpsF) * math.Max(1, math.Abs(p[j]))
		}
		copy(pp, p)
		copy(pm, p)
		pp[j] += h
		pm[j] -= h
		fp, errP := f(&Array{shape: append([]int{}, x.shape...), dt: Float, floats: pp})
		fm, errM := f(&Array{shape: append([]int{}, x.shape...), dt: Float, floats: pm})
		if errP != nil {
			return nil, base.WrapErr("Jacobian", errP)
		}
		if errM != nil {
			return nil, base.WrapErr("Jacobian", errM)
		}
		if fp.NDim() != 1 || fp.Len() != m || fm.NDim() != 1 || fm.Len() != m {
			return nil, errf("Jacobian: f must return %d values at column %d, got %d and %d", m, j, fp.Len(), fm.Len())
		}
		if fp.dt == Complex || fm.dt == Complex {
			return nil, errf("Jacobian: complex outputs are not supported")
		}
		for i := range m {
			out.floats[i*n+j] = (fp.floatAt(i) - fm.floatAt(i)) / (2 * h)
		}
	}
	return out, nil
}