// Package mathkit is a dependency-free math engine for the problems
// that actually come up when building and debugging ML code: linear
// algebra (the shape of every tensor op), calculus (gradients are how
// everything trains), probability/statistics (data, metrics, losses),
// numerical optimization (gradient descent, root finding), and the
// handful of ML-specific primitives (activations, loss functions,
// least-squares linear regression) that sit on top of those four
// pillars. It has no third-party dependency — see the comment on
// pkg/envkit/dotenv.go for why this module keeps go.mod minimal — so
// the whole engine is just the Go standard library's math/sort/fmt.
//
// Every sub-area (linalg.go, calculus.go, statistics.go, optimize.go,
// ml.go) is usable directly from Go code. solve.go adds one more
// layer on top: Solve(operation, params) dispatches a string
// operation name and a loosely-typed params map (exactly the shape
// a JSON tool-call argument object decodes to) to the right function
// and renders a human-readable report — that's the layer the agent
// service's solve_ml_math tool calls (see
// services/agent/internal/infrastructure/toolexec/mlmath.go), because
// an LLM tool call arrives as JSON, not as typed Go arguments.
package mathkit

import (
	"fmt"
	"math"
	"strings"
	"unicode"
)

// Expr is a parsed arithmetic expression that can be evaluated
// repeatedly against different variable bindings without re-parsing —
// exactly what Gradient and GradientDescent need, since they evaluate
// the same expression at many points.
type Expr struct {
	root node
	src  string
}

// String returns the original source text the expression was parsed
// from, for error messages and report rendering.
func (e *Expr) String() string { return e.src }

// node is the parsed AST. A tiny tagged union kept as an interface
// rather than a struct-with-a-kind-field: each node type carries only
// the fields it needs, and eval is a single recursive method.
type node interface {
	eval(vars map[string]float64) (float64, error)
}

type numNode float64

func (n numNode) eval(map[string]float64) (float64, error) { return float64(n), nil }

type varNode string

func (n varNode) eval(vars map[string]float64) (float64, error) {
	v, ok := vars[string(n)]
	if !ok {
		return 0, fmt.Errorf("unbound variable %q (pass it in \"variables\" or \"at\")", string(n))
	}
	return v, nil
}

type binNode struct {
	op       byte // '+', '-', '*', '/', '^'
	lhs, rhs node
}

func (n binNode) eval(vars map[string]float64) (float64, error) {
	l, err := n.lhs.eval(vars)
	if err != nil {
		return 0, err
	}
	r, err := n.rhs.eval(vars)
	if err != nil {
		return 0, err
	}
	switch n.op {
	case '+':
		return l + r, nil
	case '-':
		return l - r, nil
	case '*':
		return l * r, nil
	case '/':
		if r == 0 {
			return 0, fmt.Errorf("division by zero")
		}
		return l / r, nil
	case '^':
		return math.Pow(l, r), nil
	}
	return 0, fmt.Errorf("internal: unknown operator %q", n.op)
}

type negNode struct{ x node }

func (n negNode) eval(vars map[string]float64) (float64, error) {
	v, err := n.x.eval(vars)
	return -v, err
}

type callNode struct {
	fn   string
	args []node
}

// unaryFuncs are functions of exactly one argument — the overwhelming
// majority of what shows up in ML math (activations, logs, norms).
var unaryFuncs = map[string]func(float64) float64{
	"sin": math.Sin, "cos": math.Cos, "tan": math.Tan,
	"asin": math.Asin, "acos": math.Acos, "atan": math.Atan,
	"sinh": math.Sinh, "cosh": math.Cosh, "tanh": math.Tanh,
	"exp": math.Exp, "sqrt": math.Sqrt, "abs": math.Abs,
	"ln": math.Log, "log": math.Log, "log2": math.Log2, "log10": math.Log10,
	"floor": math.Floor, "ceil": math.Ceil, "round": math.Round,
	"sigmoid": Sigmoid, "relu": ReLU, "sign": func(x float64) float64 { return math.Copysign(1, x) * boolF(x != 0) },
}

func boolF(b bool) float64 {
	if b {
		return 1
	}
	return 0
}

// binaryFuncs are functions of exactly two arguments.
var binaryFuncs = map[string]func(float64, float64) float64{
	"max": math.Max, "min": math.Min, "pow": math.Pow,
	"mod": math.Mod, "atan2": math.Atan2,
}

func (n callNode) eval(vars map[string]float64) (float64, error) {
	switch len(n.args) {
	case 1:
		if f, ok := unaryFuncs[n.fn]; ok {
			v, err := n.args[0].eval(vars)
			if err != nil {
				return 0, err
			}
			return f(v), nil
		}
	case 2:
		if f, ok := binaryFuncs[n.fn]; ok {
			a, err := n.args[0].eval(vars)
			if err != nil {
				return 0, err
			}
			b, err := n.args[1].eval(vars)
			if err != nil {
				return 0, err
			}
			return f(a, b), nil
		}
	}
	return 0, fmt.Errorf("unknown function %q with %d argument(s)", n.fn, len(n.args))
}

// Eval evaluates the expression against a variable binding. Missing
// variables are a reported error (never silently treated as zero),
// since a silent zero would turn a typo into a wrong-but-plausible
// answer — exactly the failure mode numerical tools must not have.
func (e *Expr) Eval(vars map[string]float64) (float64, error) {
	return e.root.eval(vars)
}

// ParseExpr parses a math expression written in ordinary infix
// notation: + - * / ^ (right-associative power), unary minus,
// parentheses, named variables, the constants pi/e, and function
// calls (sin(x), sigmoid(x), max(a,b), ...). See unaryFuncs and
// binaryFuncs for the full function list.
func ParseExpr(src string) (*Expr, error) {
	p := &parser{toks: tokenize(src)}
	n, err := p.parseExpr()
	if err != nil {
		return nil, fmt.Errorf("parsing %q: %w", src, err)
	}
	if p.pos != len(p.toks) {
		return nil, fmt.Errorf("parsing %q: unexpected trailing input near %q", src, p.toks[p.pos].text)
	}
	return &Expr{root: n, src: src}, nil
}

// --- tokenizer ---

type tokKind int

const (
	tokNum tokKind = iota
	tokIdent
	tokOp
	tokLParen
	tokRParen
	tokComma
)

type token struct {
	kind tokKind
	text string
	num  float64
}

func tokenize(src string) []token {
	var toks []token
	r := []rune(src)
	i := 0
	for i < len(r) {
		c := r[i]
		switch {
		case unicode.IsSpace(c):
			i++
		case c == '(':
			toks = append(toks, token{kind: tokLParen, text: "("})
			i++
		case c == ')':
			toks = append(toks, token{kind: tokRParen, text: ")"})
			i++
		case c == ',':
			toks = append(toks, token{kind: tokComma, text: ","})
			i++
		case strings.ContainsRune("+-*/^", c):
			toks = append(toks, token{kind: tokOp, text: string(c)})
			i++
		case unicode.IsDigit(c) || c == '.':
			j := i
			for j < len(r) && (unicode.IsDigit(r[j]) || r[j] == '.') {
				j++
			}
			// scientific notation, e.g. 1e-3
			if j < len(r) && (r[j] == 'e' || r[j] == 'E') {
				k := j + 1
				if k < len(r) && (r[k] == '+' || r[k] == '-') {
					k++
				}
				if k < len(r) && unicode.IsDigit(r[k]) {
					j = k
					for j < len(r) && unicode.IsDigit(r[j]) {
						j++
					}
				}
			}
			text := string(r[i:j])
			var f float64
			fmt.Sscanf(text, "%g", &f)
			toks = append(toks, token{kind: tokNum, text: text, num: f})
			i = j
		case unicode.IsLetter(c) || c == '_':
			j := i
			for j < len(r) && (unicode.IsLetter(r[j]) || unicode.IsDigit(r[j]) || r[j] == '_') {
				j++
			}
			toks = append(toks, token{kind: tokIdent, text: string(r[i:j])})
			i = j
		default:
			// Unknown character: emit as its own operator token so the
			// parser reports a clear "unexpected token" error instead
			// of the tokenizer silently dropping it.
			toks = append(toks, token{kind: tokOp, text: string(c)})
			i++
		}
	}
	return toks
}

// --- recursive-descent parser ---
//
// expr   := term (('+' | '-') term)*
// term   := unary (('*' | '/') unary)*
// unary  := ('-' | '+') unary | power
// power  := primary ('^' unary)?         -- right-associative; the
//           exponent recurses through unary (not power) so the
//           precedence matches ordinary math/Python convention:
//           unary minus binds LOOSER than '^' (-2^2 == -(2^2) == -4)
//           but the exponent itself may carry its own sign
//           (2^-2 == 0.25).
// primary:= number | const | ident '(' args ')' | ident | '(' expr ')'
// args   := expr (',' expr)*

type parser struct {
	toks []token
	pos  int
}

func (p *parser) peek() (token, bool) {
	if p.pos >= len(p.toks) {
		return token{}, false
	}
	return p.toks[p.pos], true
}

func (p *parser) next() (token, bool) {
	t, ok := p.peek()
	if ok {
		p.pos++
	}
	return t, ok
}

func (p *parser) parseExpr() (node, error) {
	lhs, err := p.parseTerm()
	if err != nil {
		return nil, err
	}
	for {
		t, ok := p.peek()
		if !ok || t.kind != tokOp || (t.text != "+" && t.text != "-") {
			return lhs, nil
		}
		p.next()
		rhs, err := p.parseTerm()
		if err != nil {
			return nil, err
		}
		lhs = binNode{op: t.text[0], lhs: lhs, rhs: rhs}
	}
}

func (p *parser) parseTerm() (node, error) {
	lhs, err := p.parseUnary()
	if err != nil {
		return nil, err
	}
	for {
		t, ok := p.peek()
		if !ok || t.kind != tokOp || (t.text != "*" && t.text != "/") {
			return lhs, nil
		}
		p.next()
		rhs, err := p.parseUnary()
		if err != nil {
			return nil, err
		}
		lhs = binNode{op: t.text[0], lhs: lhs, rhs: rhs}
	}
}

func (p *parser) parseUnary() (node, error) {
	if t, ok := p.peek(); ok && t.kind == tokOp && (t.text == "-" || t.text == "+") {
		p.next()
		x, err := p.parseUnary()
		if err != nil {
			return nil, err
		}
		if t.text == "-" {
			return negNode{x: x}, nil
		}
		return x, nil
	}
	return p.parsePower()
}

func (p *parser) parsePower() (node, error) {
	lhs, err := p.parsePrimary()
	if err != nil {
		return nil, err
	}
	if t, ok := p.peek(); ok && t.kind == tokOp && t.text == "^" {
		p.next()
		rhs, err := p.parseUnary() // right-associative; lets the exponent carry its own sign
		if err != nil {
			return nil, err
		}
		return binNode{op: '^', lhs: lhs, rhs: rhs}, nil
	}
	return lhs, nil
}

func (p *parser) parsePrimary() (node, error) {
	t, ok := p.next()
	if !ok {
		return nil, fmt.Errorf("unexpected end of expression")
	}
	switch t.kind {
	case tokNum:
		return numNode(t.num), nil
	case tokLParen:
		n, err := p.parseExpr()
		if err != nil {
			return nil, err
		}
		if c, ok := p.next(); !ok || c.kind != tokRParen {
			return nil, fmt.Errorf("missing closing ')'")
		}
		return n, nil
	case tokIdent:
		switch strings.ToLower(t.text) {
		case "pi":
			return numNode(math.Pi), nil
		case "e":
			return numNode(math.E), nil
		}
		if nx, ok := p.peek(); ok && nx.kind == tokLParen {
			p.next()
			var args []node
			if c, ok := p.peek(); !ok || c.kind != tokRParen {
				for {
					a, err := p.parseExpr()
					if err != nil {
						return nil, err
					}
					args = append(args, a)
					c, ok := p.peek()
					if !ok {
						return nil, fmt.Errorf("missing closing ')' in call to %s", t.text)
					}
					if c.kind == tokComma {
						p.next()
						continue
					}
					break
				}
			}
			if c, ok := p.next(); !ok || c.kind != tokRParen {
				return nil, fmt.Errorf("missing closing ')' in call to %s", t.text)
			}
			return callNode{fn: strings.ToLower(t.text), args: args}, nil
		}
		return varNode(t.text), nil
	}
	return nil, fmt.Errorf("unexpected token %q", t.text)
}