"""Symbolic function abstraction built on top of ``SX`` expressions."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Iterable

from .sx import (
    SX,
    SXNode,
    SXVector,
    parse_bilinear_form_args,
    parse_matvec_component_args,
    parse_quadform_args,
    vector,
)
from .custom_elementary import (
    _invoke_custom_callback,
    evaluate_custom_hessian,
    evaluate_custom_hvp,
    evaluate_custom_jacobian,
    get_registered_elementary_function,
    parse_custom_scalar_args,
    parse_custom_scalar_hvp_args,
    parse_custom_vector_args,
    parse_custom_vector_hessian_entry_args,
    parse_custom_vector_hvp_component_args,
    parse_custom_vector_jacobian_component_args,
)


FunctionArg = SX | SXVector
BoundValue = SX | SXVector | float | int | list[object] | tuple[object, ...]


@dataclass(frozen=True, slots=True)
class Function:
    """Symbolic multi-input, multi-output function.

    A ``Function`` defines a boundary around symbolic expressions. Inputs
    must be symbolic leaves, while outputs can be arbitrary ``SX`` or
    ``SXVector`` expressions that depend on those inputs.

    The class supports:

    - explicit input and output grouping
    - validation of names and symbolic inputs
    - topological ordering of the output dependency graph
    - calling with symbolic values for substitution
    - calling with numeric values for direct evaluation
    """

    name: str
    inputs: tuple[FunctionArg, ...]
    outputs: tuple[FunctionArg, ...]
    input_names: tuple[str, ...]
    output_names: tuple[str, ...]

    def __init__(
        self,
        name: str,
        inputs: Iterable[FunctionArg],
        outputs: Iterable[FunctionArg],
        input_names: Iterable[str] | None = None,
        output_names: Iterable[str] | None = None,
    ) -> None:
        normalized_inputs = tuple(_coerce_function_arg(item, label="input") for item in inputs)
        normalized_outputs = tuple(_coerce_function_arg(item, label="output") for item in outputs)

        if not normalized_inputs:
            raise ValueError("Function must have at least one input")
        if not normalized_outputs:
            raise ValueError("Function must have at least one output")

        resolved_input_names = _resolve_names(
            names=input_names,
            count=len(normalized_inputs),
            prefix="i",
            label="input",
        )
        resolved_output_names = _resolve_names(
            names=output_names,
            count=len(normalized_outputs),
            prefix="o",
            label="output",
        )

        _validate_inputs(normalized_inputs)

        object.__setattr__(self, "name", name)
        object.__setattr__(self, "inputs", normalized_inputs)
        object.__setattr__(self, "outputs", normalized_outputs)
        object.__setattr__(self, "input_names", resolved_input_names)
        object.__setattr__(self, "output_names", resolved_output_names)

    @property
    def flat_inputs(self) -> tuple[SX, ...]:
        """Return all scalar inputs flattened into a single tuple."""
        return tuple(_flatten_args(self.inputs))

    @property
    def flat_outputs(self) -> tuple[SX, ...]:
        """Return all scalar outputs flattened into a single tuple."""
        return tuple(_flatten_args(self.outputs))

    @property
    def nodes(self) -> tuple[SXNode, ...]:
        """Return output dependency nodes in topological order."""
        ordered: list[SXNode] = []
        seen: set[SXNode] = set()

        for expr in self.flat_outputs:
            _visit_node(expr.node, seen, ordered)

        return tuple(ordered)

    def cse(self, *, prefix: str = "w", min_uses: int = 2):
        """Build a common-subexpression elimination plan for the outputs."""
        from .cse import cse

        return cse(self.outputs, prefix=prefix, min_uses=min_uses)

    def generate_rust(
        self,
        *,
        config=None,
        function_name: str | None = None,
        backend_mode: str = "std",
        scalar_type: str = "f64",
        math_library: str | None = None,
    ):
        """Generate Rust source code for primal function evaluation."""
        from .rust_codegen import generate_rust

        return generate_rust(
            self,
            config=config,
            function_name=function_name,
            backend_mode=backend_mode,
            scalar_type=scalar_type,
            math_library=math_library,
        )

    def create_rust_project(
        self,
        path: str,
        *,
        config=None,
        crate_name: str | None = None,
        function_name: str | None = None,
        backend_mode: str = "std",
        scalar_type: str = "f64",
        math_library: str | None = None,
    ):
        """Create a Rust project containing the generated primal code."""
        from .rust_codegen import create_rust_project

        return create_rust_project(
            self,
            path,
            config=config,
            crate_name=crate_name,
            function_name=function_name,
            backend_mode=backend_mode,
            scalar_type=scalar_type,
            math_library=math_library,
        )

    def create_rust_derivative_bundle(
        self,
        path: str,
        *,
        config=None,
        include_primal: bool = True,
        include_jacobians: bool = True,
        include_hessians: bool = True,
        simplify_derivatives: int | str | None = None,
    ):
        """Create a directory containing Rust projects for derivative kernels."""
        from .rust_codegen import create_rust_derivative_bundle

        return create_rust_derivative_bundle(
            self,
            path,
            config=config,
            include_primal=include_primal,
            include_jacobians=include_jacobians,
            include_hessians=include_hessians,
            simplify_derivatives=simplify_derivatives,
        )

    def __call__(self, *args: BoundValue) -> FunctionArg | tuple[FunctionArg, ...] | float | tuple[float, ...]:
        """Call the function with symbolic or numeric arguments.

        Args:
            *args: One argument for each declared input. Scalar inputs accept
                ``SX``, ``int``, or ``float``. Vector inputs accept
                ``SXVector`` or a sequence of scalar-like values.

        Returns:
            The evaluated outputs. If the call stays symbolic, the return
            value contains ``SX`` and ``SXVector`` objects. If all outputs
            reduce to numeric constants, the result is returned as ``float``
            values.
        """
        if len(args) != len(self.inputs):
            raise ValueError(
                f"expected {len(self.inputs)} arguments, received {len(args)}"
            )

        mapping: dict[SXNode, SX] = {}
        for formal, actual in zip(self.inputs, args):
            bound = _coerce_bound_arg(actual, formal)
            for formal_scalar, actual_scalar in zip(_flatten_single(formal), _flatten_single(bound)):
                mapping[formal_scalar.node] = actual_scalar

        evaluated = tuple(
            _finalize_output(_substitute_output(output, mapping)) for output in self.outputs
        )
        return _collapse_outputs(evaluated)

    def jvp(
        self,
        *tangent_inputs: BoundValue,
        name: str | None = None,
    ) -> Function:
        """Build a new function representing a forward-mode directional derivative.

        Args:
            *tangent_inputs: One tangent seed for each declared input. Seeds
                must match the shape of the corresponding inputs.
            name: Optional name for the differentiated function.

        Returns:
            A new ``Function`` with the same primal inputs and outputs equal
            to the Jacobian-vector product in the supplied tangent
            direction.
        """
        from .ad import jvp

        if len(tangent_inputs) != len(self.inputs):
            raise ValueError(
                f"expected {len(self.inputs)} tangent arguments, received {len(tangent_inputs)}"
            )

        differentiated_outputs: list[FunctionArg] = []
        for output in self.outputs:
            total = _zero_like(output)
            for wrt, tangent in zip(self.inputs, tangent_inputs):
                total = _add_like(total, jvp(output, wrt, tangent))
            differentiated_outputs.append(total)

        return Function(
            name or f"{self.name}_jvp",
            self.inputs,
            differentiated_outputs,
            input_names=self.input_names,
            output_names=self.output_names,
        )

    def gradient(self, wrt_index: int = 0, name: str | None = None) -> Function:
        """Build a gradient function for a scalar-output function."""
        if not 0 <= wrt_index < len(self.inputs):
            raise IndexError("wrt_index is out of range")
        if len(self.outputs) != 1 or not isinstance(self.outputs[0], SX):
            raise ValueError("Function.gradient requires exactly one scalar output")

        from .ad import gradient

        wrt = self.inputs[wrt_index]
        grad = gradient(self.outputs[0], wrt)
        output_name = self.output_names[0]

        if isinstance(grad, SXVector):
            outputs: list[FunctionArg] = [grad]
            output_names = (output_name,)
        else:
            outputs = [grad]
            output_names = (output_name,)

        return Function(
            name or f"{self.name}_gradient_{self.input_names[wrt_index]}",
            self.inputs,
            outputs,
            input_names=self.input_names,
            output_names=output_names,
        )

    def joint(
        self,
        components: Iterable[str],
        wrt_index: int = 0,
        name: str | None = None,
        simplify_joint: int | str | None = None,
    ) -> Function:
        """Build a function that returns several requested artifacts together.

        Supported component names are:

        - ``"f"`` for the primal outputs
        - ``"grad"`` for the gradient block with respect to ``wrt_index``
        - ``"jf"`` for the Jacobian block with respect to ``wrt_index``
        - ``"hessian"`` for the Hessian block with respect to ``wrt_index``
        - ``"hvp"`` for the Hessian-vector product with respect to ``wrt_index``

        The output order follows ``components`` exactly. If ``"hvp"`` is
        requested, the returned function appends one tangent input matching the
        selected differentiation block.
        """
        if not 0 <= wrt_index < len(self.inputs):
            raise IndexError("wrt_index is out of range")

        resolved_components = _resolve_joint_components(components)
        tangent_input: FunctionArg | None = None
        tangent_input_name: str | None = None
        if "hvp" in resolved_components:
            tangent_input_name = f"v_{self.input_names[wrt_index]}"
            tangent_input = _make_symbolic_input_like(
                self.inputs[wrt_index],
                tangent_input_name,
            )

        outputs: list[FunctionArg] = []
        output_names: list[str] = []
        for component in resolved_components:
            if component == "f":
                outputs.extend(self.outputs)
                output_names.extend(self.output_names)
                continue
            if component == "jf":
                jacobian_function = self.jacobian(wrt_index)
                outputs.extend(jacobian_function.outputs)
                output_names.extend(jacobian_function.output_names)
                continue
            if component == "grad":
                gradient_function = self.gradient(wrt_index)
                outputs.extend(gradient_function.outputs)
                output_names.extend(
                    f"gradient_{output_name}" for output_name in gradient_function.output_names
                )
                continue
            if component == "hessian":
                hessian_function = self.hessian(wrt_index)
                outputs.extend(hessian_function.outputs)
                output_names.extend(
                    f"hessian_{output_name}" for output_name in hessian_function.output_names
                )
                continue
            if component == "hvp":
                if tangent_input is None:
                    raise AssertionError("tangent input should have been created for hvp")
                if len(self.outputs) != 1 or not isinstance(self.outputs[0], SX):
                    raise ValueError("Function.joint requires exactly one scalar output when 'hvp' is requested")
                from .ad import jvp

                gradient_output = self.gradient(wrt_index).outputs[0]
                hvp_output = jvp(gradient_output, self.inputs[wrt_index], tangent_input)
                outputs.append(hvp_output)
                output_names.append(f"hvp_{self.output_names[0]}")
                continue
            raise AssertionError(f"unexpected joint component {component!r}")

        joint_name = name or _joint_function_name(
            self.name,
            self.input_names[wrt_index],
            resolved_components,
        )
        inputs = self.inputs if tangent_input is None else (*self.inputs, tangent_input)
        input_names = self.input_names if tangent_input_name is None else (*self.input_names, tangent_input_name)
        joint_function = Function(
            joint_name,
            inputs,
            outputs,
            input_names=input_names,
            output_names=output_names,
        )
        if simplify_joint is None:
            return joint_function
        return joint_function.simplify(max_effort=simplify_joint, name=joint_name)

    def hvp(
        self,
        wrt_index: int = 0,
        name: str | None = None,
        tangent_name: str | None = None,
    ) -> Function:
        """Build a Hessian-vector product function for a scalar-output function.

        The returned function keeps the original primal inputs and appends
        one additional tangent input matching the selected differentiation
        block.
        """
        if not 0 <= wrt_index < len(self.inputs):
            raise IndexError("wrt_index is out of range")
        if len(self.outputs) != 1 or not isinstance(self.outputs[0], SX):
            raise ValueError("Function.hvp requires exactly one scalar output")

        wrt = self.inputs[wrt_index]
        tangent_input_name = tangent_name or f"v_{self.input_names[wrt_index]}"
        tangent_input = _make_symbolic_input_like(wrt, tangent_input_name)
        gradient_output = self.gradient(wrt_index).outputs[0]

        from .ad import jvp

        hvp_output = jvp(gradient_output, wrt, tangent_input)

        return Function(
            name or f"{self.name}_hvp_{self.input_names[wrt_index]}",
            (*self.inputs, tangent_input),
            [hvp_output],
            input_names=(*self.input_names, tangent_input_name),
            output_names=(self.output_names[0],),
        )

    def vjp(
        self,
        *cotangent_outputs: BoundValue,
        wrt_index: int | None = None,
        name: str | None = None,
        cotangent_names: Iterable[str] | None = None,
    ) -> Function:
        """Build a new function representing a reverse-mode derivative.

        Args:
            *cotangent_outputs: One cotangent seed for each declared output.
                Seeds must match the shape of the corresponding outputs.
            wrt_index: Optional selected input block for a runtime-seeded VJP.
                When provided, the returned function appends cotangent inputs
                to the primal inputs and returns the sensitivity with respect
                to the selected input block only.
            name: Optional name for the differentiated function.
            cotangent_names: Optional names for runtime cotangent inputs.

        Returns:
            If explicit cotangent outputs are provided, returns a new
            ``Function`` with the same primal inputs and outputs equal to the
            vector-Jacobian product in the supplied cotangent direction, with
            outputs ordered like the function inputs.

            If ``wrt_index`` is provided instead, returns a function with the
            original primal inputs plus one cotangent input per declared
            output, and a single output block matching the selected input.
        """
        from .ad import vjp

        if wrt_index is not None:
            if cotangent_outputs:
                raise ValueError("runtime-seeded Function.vjp does not accept explicit cotangent outputs")
            if not 0 <= wrt_index < len(self.inputs):
                raise IndexError("wrt_index is out of range")

            resolved_cotangent_names = _resolve_cotangent_names(
                cotangent_names,
                self.output_names,
            )
            cotangent_inputs = tuple(
                _make_symbolic_input_like(output, cotangent_name)
                for output, cotangent_name in zip(self.outputs, resolved_cotangent_names)
            )

            wrt = self.inputs[wrt_index]
            total = _zero_like(wrt)
            for output, cotangent in zip(self.outputs, cotangent_inputs):
                total = _add_like(total, vjp(output, wrt, cotangent))

            return Function(
                name or f"{self.name}_vjp_{self.input_names[wrt_index]}",
                (*self.inputs, *cotangent_inputs),
                [total],
                input_names=(*self.input_names, *resolved_cotangent_names),
                output_names=(f"vjp_{self.input_names[wrt_index]}",),
            )

        if cotangent_names is not None:
            raise ValueError("cotangent_names requires wrt_index for runtime-seeded Function.vjp")

        if len(cotangent_outputs) != len(self.outputs):
            raise ValueError(
                f"expected {len(self.outputs)} cotangent arguments, received {len(cotangent_outputs)}"
            )

        differentiated_inputs: list[FunctionArg] = []
        for wrt in self.inputs:
            total = _zero_like(wrt)
            for output, cotangent in zip(self.outputs, cotangent_outputs):
                total = _add_like(total, vjp(output, wrt, cotangent))
            differentiated_inputs.append(total)

        return Function(
            name or f"{self.name}_vjp",
            self.inputs,
            differentiated_inputs,
            input_names=self.input_names,
            output_names=tuple(f"vjp_{input_name}" for input_name in self.input_names),
        )

    def jacobian(self, wrt_index: int = 0, name: str | None = None) -> Function:
        """Build a new function representing a Jacobian block.

        Args:
            wrt_index: Index of the declared input with respect to which the
                Jacobian is computed.
            name: Optional name for the differentiated function.

        Returns:
            A new ``Function`` with the same primal inputs and outputs equal
            to the Jacobian of each original output with respect to the
            selected input block.

        Notes:
            Since full matrix types are not implemented yet, vector-output
            by vector-input Jacobians are returned as one flat row-major
            ``SXVector`` per original output.
        """
        from .ad import jacobian

        if not 0 <= wrt_index < len(self.inputs):
            raise IndexError("wrt_index is out of range")

        wrt = self.inputs[wrt_index]
        differentiated_outputs: list[FunctionArg] = []
        differentiated_names: list[str] = []

        for output_name, output in zip(self.output_names, self.outputs):
            block = jacobian(output, wrt)
            differentiated_outputs.append(block)
            differentiated_names.append(f"jacobian_{output_name}")

        return Function(
            name or f"{self.name}_jacobian_{self.input_names[wrt_index]}",
            self.inputs,
            differentiated_outputs,
            input_names=self.input_names,
            output_names=differentiated_names,
        )

    def jacobian_blocks(
        self,
        wrt_indices: Iterable[int] | None = None,
    ) -> tuple[Function, ...]:
        """Build Jacobian functions for one or more input blocks."""
        indices = _resolve_block_indices(wrt_indices, len(self.inputs))
        return tuple(self.jacobian(index) for index in indices)

    def vjp_blocks(
        self,
        wrt_indices: Iterable[int] | None = None,
    ) -> tuple[Function, ...]:
        """Build runtime-seeded vector-Jacobian-product functions for input blocks."""
        indices = _resolve_block_indices(wrt_indices, len(self.inputs))
        return tuple(self.vjp(wrt_index=index) for index in indices)

    def hessian(self, wrt_index: int = 0, name: str | None = None) -> Function:
        """Build a new function representing a Hessian block.

        Args:
            wrt_index: Index of the declared input with respect to which the
                Hessian is computed.
            name: Optional name for the differentiated function.

        Returns:
            A new ``Function`` with the same primal inputs and outputs equal
            to the Hessian of the scalar output with respect to the selected
            input block.

        Notes:
            Hessians are currently defined only for single scalar-output
            functions. For vector inputs, the Hessian is returned as a
            single flat ``SXVector`` in row-major order.
        """
        from .ad import hessian

        if not 0 <= wrt_index < len(self.inputs):
            raise IndexError("wrt_index is out of range")
        if len(self.outputs) != 1 or not isinstance(self.outputs[0], SX):
            raise ValueError("Function.hessian requires exactly one scalar output")

        wrt = self.inputs[wrt_index]
        block = hessian(self.outputs[0], wrt)

        if isinstance(block, tuple):
            differentiated_outputs = [
                SXVector(tuple(entry for row in block for entry in row))
            ]
            differentiated_names = [self.output_names[0]]
        else:
            differentiated_outputs = [block]
            differentiated_names = [self.output_names[0]]

        return Function(
            name or f"{self.name}_hessian_{self.input_names[wrt_index]}",
            self.inputs,
            differentiated_outputs,
            input_names=self.input_names,
            output_names=differentiated_names,
        )

    def hessian_blocks(
        self,
        wrt_indices: Iterable[int] | None = None,
    ) -> tuple[Function, ...]:
        """Build Hessian functions for one or more input blocks."""
        indices = _resolve_block_indices(wrt_indices, len(self.inputs))
        return tuple(self.hessian(index) for index in indices)

    def hvp_blocks(
        self,
        wrt_indices: Iterable[int] | None = None,
    ) -> tuple[Function, ...]:
        """Build Hessian-vector product functions for one or more input blocks."""
        indices = _resolve_block_indices(wrt_indices, len(self.inputs))
        return tuple(self.hvp(index) for index in indices)

    def simplify(self, max_effort: int | str = "basic", name: str | None = None) -> Function:
        """Build a new function with simplified outputs.

        Args:
            max_effort: Maximum simplification effort. This can be an
                integer pass count or one of the named presets supported by
                :func:`gradgen.simplify`.
            name: Optional name for the simplified function.

        Returns:
            A new ``Function`` with the same inputs and simplified outputs.
        """
        from .simplify import _resolve_effort, _simplify_value

        effort = _resolve_effort(max_effort)
        simplified_outputs = [
            _simplify_value(output, effort) for output in self.outputs
        ]

        return Function(
            name or f"{self.name}_simplified",
            self.inputs,
            simplified_outputs,
            input_names=self.input_names,
            output_names=self.output_names,
        )

    def __repr__(self) -> str:
        return (
            f"Function(name={self.name!r}, input_names={self.input_names!r}, "
            f"output_names={self.output_names!r})"
        )


def _coerce_function_arg(item: object, *, label: str) -> FunctionArg:
    """Coerce a supported function input or output declaration."""
    if isinstance(item, (SX, SXVector)):
        return item
    raise TypeError(f"{label} must be an SX or SXVector")


def _resolve_names(
    *,
    names: Iterable[str] | None,
    count: int,
    prefix: str,
    label: str,
) -> tuple[str, ...]:
    """Resolve user-provided names or generate defaults."""
    if names is None:
        return tuple(f"{prefix}{index}" for index in range(count))

    resolved = tuple(names)
    if len(resolved) != count:
        raise ValueError(f"expected {count} {label} names, received {len(resolved)}")
    if len(set(resolved)) != len(resolved):
        raise ValueError(f"{label} names must be unique")
    return resolved


def _resolve_block_indices(
    wrt_indices: Iterable[int] | None,
    input_count: int,
) -> tuple[int, ...]:
    """Resolve and validate one or more input block indices."""
    if wrt_indices is None:
        return tuple(range(input_count))

    resolved = tuple(wrt_indices)
    for index in resolved:
        if not 0 <= index < input_count:
            raise IndexError("wrt_index is out of range")
    return resolved


def _resolve_joint_components(components: Iterable[str]) -> tuple[str, ...]:
    """Resolve and validate component names for a joint function."""
    resolved = tuple(components)
    if len(resolved) < 2:
        raise ValueError("joint functions require at least two components")
    allowed = {"f", "grad", "jf", "hessian", "hvp"}
    for component in resolved:
        if component not in allowed:
            raise ValueError(f"unsupported joint component {component!r}")
    if len(set(resolved)) != len(resolved):
        raise ValueError("joint components must be unique")
    return resolved


def _joint_function_name(
    base_name: str,
    input_name: str,
    components: tuple[str, ...],
) -> str:
    """Build the default name for a joint function."""
    return f"{base_name}_joint_{'_'.join(components)}_{input_name}"


def _resolve_cotangent_names(
    names: Iterable[str] | None,
    output_names: tuple[str, ...],
) -> tuple[str, ...]:
    """Resolve cotangent-input names associated with function outputs."""
    if names is None:
        return tuple(f"cotangent_{output_name}" for output_name in output_names)
    return _resolve_names(
        names=names,
        count=len(output_names),
        prefix="cotangent_",
        label="cotangent input",
    )


def _make_symbolic_input_like(value: FunctionArg, base_name: str) -> FunctionArg:
    """Create a fresh symbolic input with the same shape as ``value``."""
    if isinstance(value, SX):
        return SX.sym(base_name)
    return SXVector.sym(base_name, len(value))


def _validate_inputs(inputs: tuple[FunctionArg, ...]) -> None:
    """Ensure every input is built from unique symbolic leaves."""
    seen_names: set[str] = set()

    for item in inputs:
        for scalar in _flatten_single(item):
            if scalar.op != "symbol":
                raise ValueError("Function inputs must be symbolic variables")
            if scalar.name in seen_names:
                raise ValueError("Function input symbols must be unique")
            seen_names.add(scalar.name)


def _flatten_args(items: Iterable[FunctionArg]) -> Iterable[SX]:
    """Yield scalar expressions from a sequence of scalar or vector values."""
    for item in items:
        yield from _flatten_single(item)


def _flatten_single(item: FunctionArg) -> tuple[SX, ...]:
    """Flatten a single scalar or vector value into scalar expressions."""
    if isinstance(item, SX):
        return (item,)
    return item.elements


def _visit_node(node: SXNode, seen: set[SXNode], ordered: list[SXNode]) -> None:
    """Depth-first topological traversal of expression nodes."""
    if node in seen:
        return

    for arg in node.args:
        _visit_node(arg, seen, ordered)

    seen.add(node)
    ordered.append(node)


def _coerce_bound_arg(value: BoundValue, formal: FunctionArg) -> FunctionArg:
    """Coerce a call-time argument to match a declared input shape."""
    if isinstance(formal, SX):
        return _coerce_scalar_value(value)
    if isinstance(value, SXVector):
        if len(value) != len(formal):
            raise ValueError("vector argument length must match the formal input")
        return value
    if isinstance(value, (list, tuple)):
        coerced = vector(value)
        if len(coerced) != len(formal):
            raise ValueError("vector argument length must match the formal input")
        return coerced
    raise TypeError("vector inputs require an SXVector or a sequence of scalar-like values")


def _coerce_scalar_value(value: BoundValue) -> SX:
    """Coerce a scalar call-time value into an ``SX`` expression."""
    if isinstance(value, SX):
        return value
    if isinstance(value, (int, float)):
        return SX.const(value)
    raise TypeError("scalar inputs require an SX, int, or float")


def _substitute_output(output: FunctionArg, mapping: dict[SXNode, SX]) -> FunctionArg:
    """Apply input substitution to a function output."""
    if isinstance(output, SX):
        return _substitute_scalar(output, mapping)
    return SXVector(tuple(_substitute_scalar(element, mapping) for element in output))


def _substitute_scalar(expr: SX, mapping: dict[SXNode, SX]) -> SX:
    """Recursively substitute formal input symbols with bound values."""
    if expr.node in mapping:
        return mapping[expr.node]
    if expr.op in {"symbol", "const"}:
        return expr

    substituted_args = tuple(_substitute_scalar(SX(arg), mapping).node for arg in expr.node.args)
    return SX(SXNode.make(expr.op, substituted_args, name=expr.name, value=expr.value))


def _finalize_output(output: FunctionArg) -> FunctionArg | float:
    """Evaluate a substituted output when it has become fully numeric."""
    if isinstance(output, SX):
        if _is_numeric_expr(output):
            return _evaluate_scalar(output)
        return output

    if all(_is_numeric_expr(element) for element in output):
        return tuple(_evaluate_scalar(element) for element in output)
    return output


def _collapse_outputs(
    outputs: tuple[FunctionArg | float | tuple[float, ...], ...]
) -> FunctionArg | tuple[FunctionArg | float | tuple[float, ...], ...] | float | tuple[float, ...]:
    """Return a single output directly and preserve tuples for multi-output calls."""
    if len(outputs) == 1:
        return outputs[0]
    return outputs


def _zero_like(value: FunctionArg) -> FunctionArg:
    """Create a zero expression with the same shape as ``value``."""
    if isinstance(value, SX):
        return SX.const(0.0)
    return SXVector(tuple(SX.const(0.0) for _ in value))


def _add_like(left: FunctionArg, right: FunctionArg) -> FunctionArg:
    """Add scalar or vector expressions of matching shape."""
    if isinstance(left, SX) and isinstance(right, SX):
        return left + right
    if isinstance(left, SXVector) and isinstance(right, SXVector):
        return left + right
    raise TypeError("cannot add derivative values with mismatched shapes")


def _is_numeric_expr(expr: SX) -> bool:
    """Return ``True`` when an expression depends only on constants."""
    if expr.op == "const":
        return True
    if expr.op == "symbol":
        return False
    return all(_is_numeric_expr(SX(arg)) for arg in expr.node.args)


def _evaluate_scalar(expr: SX) -> float:
    """Evaluate a scalar expression made only of constants."""
    if expr.op == "const":
        if expr.value is None:
            raise ValueError("constant expression is missing a value")
        return expr.value

    args = tuple(_evaluate_scalar(SX(arg)) for arg in expr.node.args)

    if expr.op == "add":
        return args[0] + args[1]
    if expr.op == "sub":
        return args[0] - args[1]
    if expr.op == "mul":
        return args[0] * args[1]
    if expr.op == "div":
        return args[0] / args[1]
    if expr.op == "pow":
        return args[0] ** args[1]
    if expr.op == "custom_scalar":
        spec, x_value, params = parse_custom_scalar_args(expr.name, expr.args)
        if spec.eval_python is None:
            raise ValueError(f"custom function {spec.name!r} does not support numeric evaluation")
        return _invoke_custom_callback(
            spec.eval_python,
            _evaluate_scalar(x_value),
            tuple(_evaluate_scalar(param) for param in params),
            spec.parameter_dimension,
        )
    if expr.op == "custom_vector":
        spec, x_value, params = parse_custom_vector_args(expr.name, expr.args)
        if spec.eval_python is None:
            raise ValueError(f"custom function {spec.name!r} does not support numeric evaluation")
        return _invoke_custom_callback(
            spec.eval_python,
            tuple(_evaluate_scalar(value) for value in x_value),
            tuple(_evaluate_scalar(param) for param in params),
            spec.parameter_dimension,
        )
    if expr.op == "custom_scalar_jacobian":
        spec, x_value, params = parse_custom_scalar_args(expr.name, expr.args)
        return float(
            evaluate_custom_jacobian(
                spec,
                _evaluate_scalar(x_value),
                tuple(_evaluate_scalar(param) for param in params),
            )
        )
    if expr.op == "custom_scalar_hessian":
        spec, x_value, params = parse_custom_scalar_args(expr.name, expr.args)
        return float(
            evaluate_custom_hessian(
                spec,
                _evaluate_scalar(x_value),
                tuple(_evaluate_scalar(param) for param in params),
            )
        )
    if expr.op == "custom_scalar_hvp":
        spec, x_value, tangent, params = parse_custom_scalar_hvp_args(expr.name, expr.args)
        return float(
            evaluate_custom_hvp(
                spec,
                _evaluate_scalar(x_value),
                _evaluate_scalar(tangent),
                tuple(_evaluate_scalar(param) for param in params),
            )
        )
    if expr.op == "custom_vector_jacobian_component":
        spec, index, x_value, params = parse_custom_vector_jacobian_component_args(expr.name, expr.args)
        gradient = evaluate_custom_jacobian(
            spec,
            tuple(_evaluate_scalar(value) for value in x_value),
            tuple(_evaluate_scalar(param) for param in params),
        )
        if not isinstance(gradient, tuple):
            raise TypeError("vector custom Jacobians must evaluate to numeric vectors")
        return float(gradient[index])
    if expr.op == "custom_vector_hessian_entry":
        spec, row, col, x_value, params = parse_custom_vector_hessian_entry_args(expr.name, expr.args)
        hessian = evaluate_custom_hessian(
            spec,
            tuple(_evaluate_scalar(value) for value in x_value),
            tuple(_evaluate_scalar(param) for param in params),
        )
        if not isinstance(hessian, tuple):
            raise TypeError("vector custom Hessians must evaluate to numeric matrices")
        return float(hessian[row][col])
    if expr.op == "custom_vector_hvp_component":
        spec, index, x_value, tangent, params = parse_custom_vector_hvp_component_args(expr.name, expr.args)
        hvp = evaluate_custom_hvp(
            spec,
            tuple(_evaluate_scalar(value) for value in x_value),
            tuple(_evaluate_scalar(value) for value in tangent),
            tuple(_evaluate_scalar(param) for param in params),
        )
        if not isinstance(hvp, tuple):
            raise TypeError("vector custom HVPs must evaluate to numeric vectors")
        return float(hvp[index])
    if expr.op == "atan2":
        return math.atan2(args[0], args[1])
    if expr.op == "hypot":
        return math.hypot(args[0], args[1])
    if expr.op == "matvec_component":
        rows, cols, row, matrix_values, x_values = parse_matvec_component_args(expr.args)
        _ = rows
        start = row * cols
        return sum(matrix_values[start + index] * _evaluate_scalar(x_values[index]) for index in range(cols))
    if expr.op == "quadform":
        size, matrix_values, x_values = parse_quadform_args(expr.args)
        x_numeric = tuple(_evaluate_scalar(value) for value in x_values)
        total = 0.0
        for row in range(size):
            for col in range(size):
                total += matrix_values[row * size + col] * x_numeric[row] * x_numeric[col]
        return total
    if expr.op == "bilinear_form":
        rows, cols, matrix_values, x_values, y_values = parse_bilinear_form_args(expr.args)
        x_numeric = tuple(_evaluate_scalar(value) for value in x_values)
        y_numeric = tuple(_evaluate_scalar(value) for value in y_values)
        total = 0.0
        for row in range(rows):
            for col in range(cols):
                total += matrix_values[row * cols + col] * x_numeric[row] * y_numeric[col]
        return total
    if expr.op == "sum":
        return sum(args)
    if expr.op == "prod":
        total = 1.0
        for arg in args:
            total *= arg
        return total
    if expr.op == "reduce_max":
        return max(args)
    if expr.op == "reduce_min":
        return min(args)
    if expr.op == "mean":
        return sum(args) / len(args)
    if expr.op == "min":
        return min(args[0], args[1])
    if expr.op == "neg":
        return -args[0]
    if expr.op == "sin":
        return math.sin(args[0])
    if expr.op == "cos":
        return math.cos(args[0])
    if expr.op == "tan":
        return math.tan(args[0])
    if expr.op == "asin":
        return math.asin(args[0])
    if expr.op == "acos":
        return math.acos(args[0])
    if expr.op == "atan":
        return math.atan(args[0])
    if expr.op == "asinh":
        return math.asinh(args[0])
    if expr.op == "acosh":
        return math.acosh(args[0])
    if expr.op == "atanh":
        return math.atanh(args[0])
    if expr.op == "sinh":
        return math.sinh(args[0])
    if expr.op == "cosh":
        return math.cosh(args[0])
    if expr.op == "tanh":
        return math.tanh(args[0])
    if expr.op == "exp":
        return math.exp(args[0])
    if expr.op == "expm1":
        return math.expm1(args[0])
    if expr.op == "log":
        return math.log(args[0])
    if expr.op == "log1p":
        return math.log1p(args[0])
    if expr.op == "sqrt":
        return math.sqrt(args[0])
    if expr.op == "cbrt":
        return math.copysign(abs(args[0]) ** (1.0 / 3.0), args[0])
    if expr.op == "erf":
        return math.erf(args[0])
    if expr.op == "erfc":
        return math.erfc(args[0])
    if expr.op == "floor":
        return math.floor(args[0])
    if expr.op == "ceil":
        return math.ceil(args[0])
    if expr.op == "round":
        if args[0] >= 0.0:
            return math.floor(args[0] + 0.5)
        return math.ceil(args[0] - 0.5)
    if expr.op == "trunc":
        return math.trunc(args[0])
    if expr.op == "fract":
        return args[0] - math.trunc(args[0])
    if expr.op == "signum":
        if args[0] > 0.0:
            return 1.0
        if args[0] < 0.0:
            return -1.0
        return 0.0
    if expr.op == "norm2":
        return math.sqrt(sum(arg * arg for arg in args))
    if expr.op == "norm2sq":
        return sum(arg * arg for arg in args)
    if expr.op == "norm1":
        return sum(math.fabs(arg) for arg in args)
    if expr.op == "norm_inf":
        return max(math.fabs(arg) for arg in args)
    if expr.op == "norm_p_to_p":
        p = args[-1]
        return sum(math.fabs(arg) ** p for arg in args[:-1])
    if expr.op == "norm_p":
        p = args[-1]
        return sum(math.fabs(arg) ** p for arg in args[:-1]) ** (1.0 / p)
    if expr.op == "abs":
        return math.fabs(args[0])
    if expr.op == "max":
        return max(args[0], args[1])

    raise ValueError(f"cannot evaluate operation {expr.op!r}")
