package mathkit

import (
	"fmt"
	"math"
)

// GradientDescentResult is the trace of one gradient-descent run.
type GradientDescentResult struct {
	// Path holds the point visited at the start of every iteration,
	// Path[0] being the caller's starting point and the last entry
	// being the final point — so len(Path) == Iterations+1. Kept in
	// full (rather than just the final point) because seeing the
	// trajectory is usually the point of running this: it's how you
	// tell "converging nicely" apart from "diverging" or
	// "oscillating because the learning rate is too high".
	Path [][]float64
	// Values holds f(Path[i]) for each point in Path, same length.
	Values []float64
	// Converged is true if the gradient norm dropped below the
	// requested tolerance before Iterations ran out.
	Converged bool
}

// GradientDescent runs plain (batch) gradient descent on expr, over
// the variables listed (in order — that order defines the vector
// space the descent happens in), starting at `start`, for at most
// `iterations` steps of size `learningRate`. tol<=0 disables the
// early-stop check (it always runs the full iteration count).
//
// This is deliberately vanilla gradient descent — no momentum, no
// Adam-style adaptive rates — because the point of this tool is to
// let a user see and reason about the bare update rule
// x := x - lr * ∇f(x) that every fancier optimizer builds on, not to
// re-implement a training framework.
func GradientDescent(expr *Expr, variables []string, start []float64, learningRate float64, iterations int, tol float64) (*GradientDescentResult, error) {
	if len(variables) != len(start) {
		return nil, fmt.Errorf("gradient_descent: %d variable name(s) but start point has %d coordinate(s)", len(variables), len(start))
	}
	if learningRate <= 0 {
		return nil, fmt.Errorf("gradient_descent: learning_rate must be positive")
	}
	if iterations <= 0 {
		iterations = 100
	}
	point := append([]float64(nil), start...)
	res := &GradientDescentResult{}
	addPoint := func(p []float64) error {
		vars := make(map[string]float64, len(variables))
		for i, name := range variables {
			vars[name] = p[i]
		}
		v, err := expr.Eval(vars)
		if err != nil {
			return err
		}
		res.Path = append(res.Path, append([]float64(nil), p...))
		res.Values = append(res.Values, v)
		return nil
	}
	if err := addPoint(point); err != nil {
		return nil, fmt.Errorf("gradient_descent: evaluating start point: %w", err)
	}
	vars := make(map[string]float64, len(variables))
	for iter := 0; iter < iterations; iter++ {
		for i, name := range variables {
			vars[name] = point[i]
		}
		grad, err := Gradient(expr, variables, vars, 0)
		if err != nil {
			return nil, fmt.Errorf("gradient_descent: iteration %d: %w", iter, err)
		}
		if tol > 0 && Norm(grad, 2) < tol {
			res.Converged = true
			break
		}
		for i := range point {
			point[i] -= learningRate * grad[i]
		}
		if err := addPoint(point); err != nil {
			return nil, fmt.Errorf("gradient_descent: iteration %d: %w", iter, err)
		}
		if math.IsNaN(point[0]) || math.IsInf(point[0], 0) {
			return nil, fmt.Errorf("gradient_descent: diverged (values went to NaN/Inf) — try a smaller learning_rate")
		}
	}
	return res, nil
}

// NewtonRootResult is the trace of one Newton's-method root-finding
// run on a single-variable expression.
type NewtonRootResult struct {
	Path      []float64 // x_0, x_1, ..., x_final
	Converged bool
}

// NewtonRoot finds a root of expr(variable) = 0 via Newton's method:
// x_{n+1} = x_n - f(x_n)/f'(x_n), with f' estimated by central
// difference (see Derivative). Any other free variables in expr must
// be supplied in `fixed`. tol<=0 defaults to 1e-10.
func NewtonRoot(expr *Expr, variable string, start float64, iterations int, tol float64, fixed map[string]float64) (*NewtonRootResult, error) {
	if iterations <= 0 {
		iterations = 100
	}
	if tol <= 0 {
		tol = 1e-10
	}
	x := start
	res := &NewtonRootResult{Path: []float64{x}}
	vars := cloneVars(fixed)
	for iter := 0; iter < iterations; iter++ {
		vars[variable] = x
		fx, err := expr.Eval(vars)
		if err != nil {
			return nil, fmt.Errorf("newton_root: iteration %d: %w", iter, err)
		}
		if math.Abs(fx) < tol {
			res.Converged = true
			break
		}
		dfx, err := Derivative(expr, variable, vars, 0)
		if err != nil {
			return nil, fmt.Errorf("newton_root: iteration %d: %w", iter, err)
		}
		if math.Abs(dfx) < 1e-14 {
			return nil, fmt.Errorf("newton_root: derivative vanished near x=%g (flat spot) — try a different start", x)
		}
		x = x - fx/dfx
		res.Path = append(res.Path, x)
		if math.IsNaN(x) || math.IsInf(x, 0) {
			return nil, fmt.Errorf("newton_root: diverged (x went to NaN/Inf) — try a different start")
		}
	}
	return res, nil
}