package mathkit

import "fmt"

// defaultStep is the finite-difference step size used when the
// caller doesn't specify one. 1e-5 is the standard sweet spot for
// float64 central differences: small enough for truncation error to
// be negligible, large enough that it doesn't get swallowed by
// float64 rounding error (the classic bias-variance trade-off of
// numerical differentiation).
const defaultStep = 1e-5

// Derivative estimates d/d(variable) of expr at the point `at`
// (variable -> value, including any other free variables the
// expression needs) using the central-difference formula
// (f(x+h)-f(x-h))/(2h), which is second-order accurate — far better
// than the naive forward difference (f(x+h)-f(x))/h for the same h.
// h<=0 uses defaultStep.
func Derivative(expr *Expr, variable string, at map[string]float64, h float64) (float64, error) {
	if h <= 0 {
		h = defaultStep
	}
	x0, ok := at[variable]
	if !ok {
		return 0, fmt.Errorf("derivative: %q is not bound in \"at\"", variable)
	}
	plus := cloneVars(at)
	plus[variable] = x0 + h
	minus := cloneVars(at)
	minus[variable] = x0 - h
	fPlus, err := expr.Eval(plus)
	if err != nil {
		return 0, fmt.Errorf("derivative: evaluating at %s+h: %w", variable, err)
	}
	fMinus, err := expr.Eval(minus)
	if err != nil {
		return 0, fmt.Errorf("derivative: evaluating at %s-h: %w", variable, err)
	}
	return (fPlus - fMinus) / (2 * h), nil
}

// Gradient estimates ∇f at the point `at` with respect to each name
// in variables, in order — the central operation every ML training
// loop runs millions of times (just normally via backprop's chain
// rule rather than finite differences; this is the numerical
// stand-in that's useful for checking a hand-derived gradient, or for
// expressions where deriving the symbolic gradient isn't worth it).
func Gradient(expr *Expr, variables []string, at map[string]float64, h float64) ([]float64, error) {
	if len(variables) == 0 {
		return nil, fmt.Errorf("gradient: \"variables\" must list at least one variable")
	}
	grad := make([]float64, len(variables))
	for i, v := range variables {
		d, err := Derivative(expr, v, at, h)
		if err != nil {
			return nil, err
		}
		grad[i] = d
	}
	return grad, nil
}

func cloneVars(m map[string]float64) map[string]float64 {
	out := make(map[string]float64, len(m))
	for k, v := range m {
		out[k] = v
	}
	return out
}

// Integrate estimates the definite integral of expr with respect to
// variable over [from, to] using Simpson's rule, which is exact for
// up to cubic polynomials and converges much faster than the
// trapezoid rule for smooth functions. steps (the number of
// subintervals) is rounded up to the nearest even number, since
// Simpson's rule needs an even number of subintervals; steps<=0
// defaults to 1000.
func Integrate(expr *Expr, variable string, from, to float64, steps int, fixed map[string]float64) (float64, error) {
	if steps <= 0 {
		steps = 1000
	}
	if steps%2 != 0 {
		steps++
	}
	h := (to - from) / float64(steps)
	vars := cloneVars(fixed)

	eval := func(x float64) (float64, error) {
		vars[variable] = x
		return expr.Eval(vars)
	}

	f0, err := eval(from)
	if err != nil {
		return 0, fmt.Errorf("integral: %w", err)
	}
	fn, err := eval(to)
	if err != nil {
		return 0, fmt.Errorf("integral: %w", err)
	}
	sum := f0 + fn
	for i := 1; i < steps; i++ {
		x := from + float64(i)*h
		fx, err := eval(x)
		if err != nil {
			return 0, fmt.Errorf("integral: %w", err)
		}
		if i%2 == 0 {
			sum += 2 * fx
		} else {
			sum += 4 * fx
		}
	}
	return sum * h / 3, nil
}