import logging
from collections.abc import Sequence
from copy import copy
from functools import singledispatch

import numpy as np

import pytensor.scalar as ps
from pytensor.compile.inner_function import HasInnerFunction
from pytensor.gradient import DisconnectedType, grad, jacobian
from pytensor.graph.basic import Apply, Constant
from pytensor.graph.fg import FrozenFunctionGraph, FunctionGraph
from pytensor.graph.null_type import NullType
from pytensor.graph.op import (
    ComputeMapType,
    Op,
    StorageMapType,
    io_connection_pattern,
)
from pytensor.graph.replace import graph_replace
from pytensor.graph.traversal import (
    ancestors,
    graph_inputs,
    truncated_graph_inputs,
)
from pytensor.scalar import ScalarType, ScalarVariable
from pytensor.tensor import as_tensor_variable
from pytensor.tensor.basic import (
    atleast_2d,
    scalar_from_tensor,
    tensor,
    tensor_from_scalar,
    zeros_like,
)
from pytensor.tensor.linalg.solvers.general import solve
from pytensor.tensor.math import tensordot
from pytensor.tensor.reshape import pack, unpack
from pytensor.tensor.variable import TensorVariable, Variable


# scipy.optimize can be slow to import, and will not be used by most users
# We import scipy.optimize lazily inside optimization perform methods to avoid this.
optimize = None


_log = logging.getLogger(__name__)


class LRUCache1:
    """
    Simple LRU cache with a memory size of 1.

    This cache is only usable for a function that takes a single input `x` and returns a single output. The
    function can also take any number of additional arguments `*args`, but these are assumed to be constant
    between function calls.

    The purpose of this cache is to allow for Hessian computation to be reused when calling scipy.optimize functions.
    It is very often the case that some sub-computations are repeated between the objective, gradient, and hessian
    functions, but by default scipy only allows for the objective and gradient to be fused.

    By using this cache, all 3 functions can be fused, which can significantly speed up the optimization process for
    expensive functions.
    """

    def __init__(self, fn, copy_x: bool = False):
        self.fn = fn
        self.last_x = None
        self.last_result = None
        self.copy_x = copy_x

        # Scipy does not respect dtypes *at all*, so we have to force it ourselves.
        self.dtype = fn.maker.fgraph.inputs[0].type.dtype

        self.cache_hits = 0
        self.cache_misses = 0

        self.value_calls = 0
        self.grad_calls = 0
        self.value_and_grad_calls = 0
        self.hess_calls = 0

    def __call__(self, x, *args):
        """
        Call the cached function with the given input `x` and additional arguments `*args`.

        If the input `x` is the same as the last input, return the cached result. Otherwise update the cache with the
        new input and result.
        """
        x = x.astype(self.dtype)

        if self.last_result is None or not (x == self.last_x).all():
            self.cache_misses += 1

            # scipy.optimize.root changes x in place, so the cache has to copy it, otherwise we get false
            # cache hits and optimization always fails.
            if self.copy_x:
                x = x.copy()
            self.last_x = x

            result = self.fn(x, *args)
            self.last_result = result

            return result

        else:
            self.cache_hits += 1
            return self.last_result

    def value(self, x, *args):
        self.value_calls += 1
        return self(x, *args)[0]

    def grad(self, x, *args):
        self.grad_calls += 1
        return self(x, *args)[1]

    def value_and_grad(self, x, *args):
        self.value_and_grad_calls += 1
        return self(x, *args)[:2]

    def hess(self, x, *args):
        self.hess_calls += 1
        return self(x, *args)[-1]

    def report(self):
        _log.info(f"Value and Grad calls: {self.value_and_grad_calls}")
        _log.info(f"Hess Calls: {self.hess_calls}")
        _log.info(f"Hits: {self.cache_hits}")
        _log.info(f"Misses: {self.cache_misses}")

    def clear_cache(self):
        self.last_x = None
        self.last_result = None
        self.cache_hits = 0
        self.cache_misses = 0
        self.value_calls = 0
        self.grad_calls = 0
        self.value_and_grad_calls = 0
        self.hess_calls = 0


def _find_optimization_parameters(
    objective: TensorVariable, x: TensorVariable
) -> list[Variable]:
    """
    Find the parameters of the optimization problem that are not the variable `x`.

    This is used to determine the additional arguments that need to be passed to the objective function.
    """

    def _depends_only_on_constants(var: Variable) -> bool:
        if isinstance(var, Constant):
            return True
        if var.owner is None:
            return False
        return all(isinstance(v, Constant) for v in graph_inputs([var]))

    return [
        arg
        for arg in truncated_graph_inputs([objective], [x])
        if (arg is not x and not _depends_only_on_constants(arg))
    ]


def _tensorize_scalar_parameters(
    objective: TensorVariable, args: list[Variable]
) -> tuple[TensorVariable, list[Variable], list[Variable]]:
    """Replace ``ScalarType`` parameters with equivalent ``TensorType`` inputs.

    scipy hands every value to the objective as a numpy array, and stores/forwards
    a ``ScalarType`` value as a bare python float -- which no C thunk accepts under
    ``trust_input``. Re-expressing such a parameter as a 0-d ``TensorType`` input
    (reconstructed inside the graph with ``scalar_from_tensor``) makes the op
    boundary array-typed, so scipy passes a preserved 0-d ndarray. Returns the
    rewritten objective, the inner-graph parameters, and the matching outer inputs.
    """
    replacements: dict[Variable, Variable] = {}
    inner_args: list[Variable] = []
    outer_args: list[Variable] = []
    for arg in args:
        if isinstance(arg.type, ScalarType):
            inner = tensor(name=arg.name, shape=(), dtype=arg.type.dtype)
            replacements[arg] = scalar_from_tensor(inner)
            inner_args.append(inner)
            outer_args.append(tensor_from_scalar(arg))
        else:
            inner_args.append(arg)
            outer_args.append(arg)

    if replacements:
        [objective] = graph_replace([objective], replacements)
    return objective, inner_args, outer_args


class ScipyWrapperOp(HasInnerFunction, Op):
    """Shared logic for scipy optimization ops.

    The inner graph is held frozen (immutable) as ``self.fgraph``, so the
    canonical op can never be mutated in place. Graph manipulation (``grad`` /
    ``graph_replace`` / ``function``) runs on a fresh mutable copy obtained via
    ``self.fgraph.unfreeze()`` -- ``graph_replace`` on frozen nodes would
    otherwise leak them into the outer graph.
    """

    # Attribute names (besides the frozen inner graph) that distinguish two ops
    # of the same type for eq/hash. Subclasses override.
    _scipy_props: tuple[str, ...] = ()

    @property
    def fn_wrapped(self):
        if self._fn_wrapped is None:
            self._fn_wrapped = LRUCache1(self.fn)
        return self._fn_wrapped

    def _prop_values(self):
        return tuple(getattr(self, name) for name in self._scipy_props)

    def __eq__(self, other):
        if self is other:
            return True
        if type(self) is not type(other):
            return False
        return (
            self.fgraph == other.fgraph and self._prop_values() == other._prop_values()
        )

    def __hash__(self):
        # ``optimizer_kwargs`` may hold nested dicts, so the props stay out of
        # the hash; equal ops still hash equal on the type and inner graph.
        return hash((type(self), self.fgraph))

    def clone_with_inner_graph(self, inner_fgraph):
        """Return a copy of this op whose inner graph is ``inner_fgraph``.

        Used by the ``optimize_inner_graph`` rewrite to bake an
        already-optimized inner graph into a NEW immutable op without touching
        ``self``. ``inner_fgraph`` may be a mutable ``FunctionGraph`` (it is
        frozen here) or an already-frozen graph.
        """
        clone_op = copy(self)
        clone_op._fn = None
        clone_op._fn_wrapped = None
        clone_op.fgraph = (
            inner_fgraph
            if isinstance(inner_fgraph, FrozenFunctionGraph)
            else inner_fgraph.freeze()
        )
        return clone_op

    # Name used by the canonicalization rewrite that rebuilds the inner graph.
    clone_with_new_fgraph = clone_with_inner_graph

    def clone(self):
        # The inner graph is immutable (a frozen ``FunctionGraph``), so there is
        # nothing to deep-clone -- mirror ``Composite``/``OpFromGraph``.
        return self

    def prepare_node(
        self,
        node: Apply,
        storage_map: StorageMapType | None,
        compute_map: ComputeMapType | None,
        impl: str | None,
    ):
        """Trigger the compilation of the inner fgraph so it shows in the dprint before the first call"""
        self.fn

    def make_node(self, *inputs):
        assert len(inputs) == len(self.inner_inputs)
        inputs = [
            inner_input.type.filter_variable(input)
            for input, inner_input in zip(inputs, self.inner_inputs)
        ]

        return Apply(self, inputs, [self.inner_inputs[0].type(), ps.bool("success")])


@singledispatch
def rewrite_optimize_inner_graph(linker, op, node, inner, *, mode):
    """Rewrite an ``optimize`` op (Minimize/Root) inner graph for ``linker``.

    Each linker registers its own; see ``pytensor.tensor.rewriting.optimize``.
    """
    raise NotImplementedError(
        f"Linker {type(linker).__name__} has not registered an optimize-op "
        "inner-graph rewrite"
    )


class ScipyScalarWrapperOp(ScipyWrapperOp):
    def compile_fn(self, mode):
        # We need to adjust the graph to work with what scipy will be passing into the inner function --
        # always scalar array of float64 type
        fgraph = self.fgraph.unfreeze()
        x, *args = fgraph.inputs
        new_root_x = ps.float64(name="x_scalar")
        new_x = tensor_from_scalar(new_root_x.astype(x.type.dtype))

        new_outputs = graph_replace(fgraph.outputs, {x: new_x})

        return self.link_fgraph(
            FunctionGraph([new_root_x, *args], new_outputs, clone=False), mode
        )

    def compute_implicit_gradients(
        self,
        x_star: TensorVariable,
        args: Sequence[TensorVariable | ScalarVariable],
        output_grad: TensorVariable,
        is_minimization: bool,
    ):
        """
        Compute gradients of a scalar optimization problem with respect to its parameters.

        For details, see the docstring of ScipyVectorWrapperOp.compute_implicit_gradients.

        Parameters
        ----------
        x_star : TensorVariable
            The symbolic solution of the optimization problem.
        args : Sequence of TensorVariable or ScalarVariable
            The parameters of the optimization problem.
        output_grad : TensorVariable
            The gradient of the output of the optimization Op with respect to some scalar loss.
        is_minimization : bool
            Whether the optimization problem is a minimization problem. If False, it is assumed to be a root-finding
            problem.
        """
        fgraph = self.fgraph.unfreeze()
        inner_x, *inner_args = fgraph.inputs
        inner_fx = fgraph.outputs[0]

        if is_minimization:
            # The implicit function in minimization is grad(x, theta) == 0
            inner_fx = grad(inner_fx, inner_x)

        df_dx, *arg_grads = grad(
            inner_fx,
            [inner_x, *inner_args],
            disconnected_inputs="ignore",
            null_gradients="return",
            return_disconnected="disconnected",
        )

        args_to_diff: tuple[bool, ...] = tuple(
            not isinstance(g.type, DisconnectedType | NullType) for g in arg_grads
        )

        if not any(args_to_diff):
            # No differentiable arguments, return disconnected gradients
            return arg_grads

        df_dthetas = [g for g, to_diff in zip(arg_grads, args_to_diff) if to_diff]

        # Make gradient an expression of the outer variables
        df_dx_star, *df_dthetas_stars = graph_replace(
            [df_dx, *df_dthetas], replace=tuple(zip(fgraph.inputs, (x_star, *args)))
        )

        grad_wrt_args = []
        df_dthetas_iter = iter(df_dthetas_stars)
        for i, (arg, to_diff) in enumerate(zip(args, args_to_diff)):
            if not to_diff:
                # Store the null grad we got from the initial `grad` call
                g = arg_grads[i]
                assert isinstance(g.type, NullType | DisconnectedType)
            else:
                # Compute non-null grad and chain with output_grad
                df_dtheta_star = next(df_dthetas_iter)
                g = (-df_dtheta_star / df_dx_star) * output_grad
            grad_wrt_args.append(g)

        assert next(df_dthetas_iter, None) is None, "Iterator was not exhausted"

        return grad_wrt_args


class ScipyVectorWrapperOp(ScipyWrapperOp):
    def compile_fn(self, mode):
        # We need to adjust the graph to work with what scipy will be passing into the inner function --
        # always a vector array with size of at least 1
        if self.inner_inputs[0].type.shape != ():
            return super().compile_fn(mode)

        fgraph = self.fgraph.unfreeze()
        x, *args = fgraph.inputs
        new_root_x = x[None].type()
        new_x = new_root_x.squeeze()
        new_outputs = graph_replace(fgraph.outputs, {x: new_x})
        return self.link_fgraph(
            FunctionGraph([new_root_x, *args], new_outputs, clone=False), mode
        )

    def compute_implicit_gradients(
        self,
        x_star: TensorVariable,
        args: Sequence[TensorVariable | ScalarVariable],
        output_grad: TensorVariable,
        is_minimization: bool,
    ):
        r"""
        Compute gradients of an optimization problem with respect to its parameters.

        Parameters
        ----------
        x_star : TensorVariable
            The symbolic solution of the optimization problem.
        args : Sequence of TensorVariable or ScalarVariable
            The parameters of the optimization problem.
        output_grad : TensorVariable
            The gradient of the output of the optimization Op with respect to some scalar loss.
        is_minimization : bool
            Whether the optimization problem is a minimization problem. If False, it is assumed to be a root-finding
            problem.

        Notes
        -----
        The gradients are computed using the implicit function theorem. Given a function `f(x, theta) = 0`, and a function
        `x_star(theta)` that, given input parameters theta returns `x` such that `f(x_star(theta), theta) = 0`, we can
        compute the gradients of `x_star` with respect to `theta` as follows:

        .. math::

            \underbrace{\frac{\partial f}{\partial x}\left(x^*(\theta), \theta\right)}_{\text{Jacobian wrt } x}
            \frac{d x^*(\theta)}{d \theta}
            +
            \underbrace{\frac{\partial f}{\partial \theta}\left(x^*(\theta), \theta\right)}_{\text{Jacobian wrt } \theta}
            = 0

        Which, after rearranging, gives us:

        .. math::

            \frac{d x^*(\theta)}{d \theta} = - \left(\frac{\partial f}{\partial x}\left(x^*(\theta),
            \theta\right)\right)^{-1} \frac{\partial f}{\partial \theta}\left(x^*(\theta), \theta\right)

        Note that this method assumes `f(x_star(theta), theta) = 0`; so it is not immediately applicable to a minimization
        problem, where `f` is the objective function. In this case, we instead take `f` to be the gradient of the objective
        function, which *is* indeed zero at the minimum.
        """
        fgraph = self.fgraph.unfreeze()
        inner_x, *inner_args = fgraph.inputs
        implicit_f = fgraph.outputs[0]
        if is_minimization:
            # The implicit function in minimization is grad(x, theta) == 0
            implicit_f = grad(implicit_f, inner_x)

        # Call grad to see what arguments are connected
        _, *arg_grads = grad(
            implicit_f.sum(),
            [inner_x, *inner_args],
            disconnected_inputs="ignore",
            null_gradients="return",
            return_disconnected="disconnected",
        )

        args_to_diff: tuple[bool, ...] = tuple(
            not isinstance(g.type, DisconnectedType | NullType) for g in arg_grads
        )

        if not args_to_diff:
            # No differentiable arguments, return disconnected/null gradients
            return arg_grads

        # Gradients are computed using the inner graph of the optimization op, not the actual inputs/outputs of the op.
        packed_inner_args, packed_arg_shapes, implicit_f = pack_inputs_of_objective(
            implicit_f,
            [
                inner_arg
                for inner_arg, to_diff in zip(inner_args, args_to_diff)
                if to_diff
            ],
        )

        df_dx, df_dtheta = jacobian(
            implicit_f,
            [inner_x, packed_inner_args],
            disconnected_inputs="ignore",
            vectorize=self.use_vectorized_jac,
        )

        # Replace inner inputs (abstract dummies) with outer inputs (the actual user-provided symbols)
        # at the solution point. Innner arguments aren't needed anymore, delete them to avoid accidental references.
        del inner_x, inner_args
        inner_to_outer_map = tuple(zip(fgraph.inputs, (x_star, *args)))
        df_dx_star, df_dtheta_star = graph_replace(
            [df_dx, df_dtheta],
            replace=inner_to_outer_map,
        )

        if df_dtheta_star.ndim == 0 or df_dx_star.ndim == 0:
            grad_wrt_args_packed = -(df_dtheta_star / df_dx_star)
        else:
            grad_wrt_args_packed = solve(-atleast_2d(df_dx_star), df_dtheta_star)

        if packed_arg_shapes is not None:
            packed_shapes_from_outer = graph_replace(
                packed_arg_shapes, inner_to_outer_map, strict=False
            )
            grad_wrt_args = unpack(
                grad_wrt_args_packed,
                packed_shapes_from_outer,
                keep_axes=None if all(inp.ndim == 0 for inp in (x_star, *args)) else 0,
            )
        else:
            grad_wrt_args = [grad_wrt_args_packed]

        final_grads = []
        grad_wrt_args_iter = iter(grad_wrt_args)
        for i, (arg, to_diff) in enumerate(zip(args, args_to_diff)):
            if not to_diff:
                # Store the null grad we got from the initial `grad` call
                g = arg_grads[i]
                assert isinstance(g.type, NullType | DisconnectedType)
            else:
                # Compute non-null grad and chain with output_grad
                arg_grad = next(grad_wrt_args_iter)
                if arg_grad.ndim > 0 and output_grad.ndim > 0:
                    g = tensordot(output_grad, arg_grad, [[0], [0]])
                else:
                    g = arg_grad * output_grad
                if isinstance(arg.type, ScalarType) and isinstance(g, TensorVariable):
                    g = scalar_from_tensor(g)
            final_grads.append(g)

        assert next(grad_wrt_args_iter, None) is None, "Iterator was not exhausted"

        return final_grads


def _optimizer_connection_pattern(fgraph, is_minimization):
    """Compute connection pattern for scipy optimization Ops.

    There are two outputs: the optimized variable, and a success flag. The success flag is
    not differentiable, so it is never connected. The optimized variable is connected only
    to inputs that are connected to the implicit function used in the gradient computation.

    For minimization, the implicit function is grad(objective, x), not the objective itself.
    An input may be connected to the objective but disconnected from its gradient (e.g. an
    additive constant), so the connection pattern must reflect the actual implicit function.
    """
    fgraph = fgraph.unfreeze()
    inner_x = fgraph.inputs[0]
    fx = fgraph.outputs[0]
    if is_minimization:
        fx = grad(fx, inner_x)
    return [
        [connected, False] for [connected] in io_connection_pattern(fgraph.inputs, [fx])
    ]


class MinimizeScalarOp(ScipyScalarWrapperOp):
    _scipy_props = ("method", "optimizer_kwargs")

    def __init__(
        self,
        x: TensorVariable,
        *args: Variable,
        objective: TensorVariable,
        method: str,
        optimizer_kwargs: dict | None = None,
    ):
        if not (isinstance(x, TensorVariable) and x.ndim == 0):
            raise ValueError(
                "The variable `x` must be a scalar (0-dimensional) tensor for minimize_scalar."
            )
        if not (isinstance(objective, TensorVariable) and objective.ndim == 0):
            raise ValueError(
                "The objective function must be a scalar (0-dimensional) tensor for minimize_scalar."
            )
        if x not in ancestors([objective]):
            raise ValueError(
                "The variable `x` must be an input to the computational graph of the objective function."
            )
        self.fgraph = FrozenFunctionGraph.from_io([x, *args], [objective])

        self.method = method
        self.optimizer_kwargs = optimizer_kwargs if optimizer_kwargs is not None else {}
        self._fn = None
        self._fn_wrapped = None

    def connection_pattern(self, node=None):
        return _optimizer_connection_pattern(self.fgraph, is_minimization=True)

    def __str__(self):
        return f"{self.__class__.__name__}(method={self.method})"

    def perform(self, node, inputs, outputs):
        global optimize
        if optimize is None:
            import scipy.optimize as optimize

        f = self.fn_wrapped
        f.clear_cache()

        # minimize_scalar doesn't take x0 as an argument. The Op still needs this input (to symbolically determine
        # the args of the objective function), but it is not used in the optimization.
        x0, *args = inputs

        res = optimize.minimize_scalar(
            fun=f.value,
            args=tuple(args),
            method=self.method,
            **self.optimizer_kwargs,
        )

        outputs[0][0] = np.array(res.x, dtype=x0.dtype)
        outputs[1][0] = np.bool_(res.success)

    def pullback(self, inputs, outputs, output_grads):
        # TODO: Handle disconnected inputs
        x, *args = inputs
        x_star, _ = outputs
        output_grad, _ = output_grads

        d_xstar_d_theta = self.compute_implicit_gradients(
            x_star=x_star,
            args=args,
            output_grad=output_grad,
            is_minimization=True,
        )

        return [zeros_like(x), *d_xstar_d_theta]


def minimize_scalar(
    objective: TensorVariable,
    x: TensorVariable,
    method: str = "brent",
    optimizer_kwargs: dict | None = None,
):
    """
    Minimize a scalar objective function using scipy.optimize.minimize_scalar.

    Parameters
    ----------
    objective : TensorVariable
        The objective function to minimize. This should be a PyTensor variable representing a scalar value.
    x : TensorVariable
        The variable with respect to which the objective function is minimized. It must be a scalar and an
        input to the computational graph of `objective`.
    method : str, optional
        The optimization method to use. Default is "brent". See `scipy.optimize.minimize_scalar` for other options.
    optimizer_kwargs : dict, optional
        Additional keyword arguments to pass to `scipy.optimize.minimize_scalar`.

    Returns
    -------
    solution: TensorVariable
        Value of `x` that minimizes `objective(x, *args)`. If the success flag is False, this will be the
        final state returned by the minimization routine, not necessarily a minimum.
    success : TensorVariable
        Symbolic boolean flag indicating whether the minimization routine reported convergence to a minimum
        value, based on the requested convergence criteria.
    """
    args = _find_optimization_parameters(objective, x)
    objective, args, outer_args = _tensorize_scalar_parameters(objective, args)

    minimize_scalar_op = MinimizeScalarOp(
        x,
        *args,
        objective=objective,
        method=method,
        optimizer_kwargs=optimizer_kwargs,
    )

    solution, success = minimize_scalar_op(x, *outer_args)

    return solution, success


class MinimizeOp(ScipyVectorWrapperOp):
    _scipy_props = (
        "method",
        "jac",
        "hess",
        "hessp",
        "use_vectorized_jac",
        "optimizer_kwargs",
    )

    def __init__(
        self,
        x: TensorVariable,
        *args: Variable,
        objective: TensorVariable,
        method: str,
        jac: bool = True,
        hess: bool = False,
        hessp: bool = False,
        use_vectorized_jac: bool = False,
        optimizer_kwargs: dict | None = None,
    ):
        if not (isinstance(x, TensorVariable) and x.ndim in (0, 1)):
            raise ValueError(
                "The variable `x` must be a scalar or vector (0-or-1-dimensional) tensor for minimize."
            )
        if not (isinstance(objective, TensorVariable) and objective.ndim == 0):
            raise ValueError(
                "The objective function must be a scalar (0-dimensional) tensor for minimize."
            )
        if x not in ancestors([objective]):
            raise ValueError(
                "The variable `x` must be an input to the computational graph of the objective function."
            )

        self.fgraph = FunctionGraph([x, *args], [objective])
        self.use_vectorized_jac = use_vectorized_jac

        if jac:
            grad_wrt_x = grad(self.fgraph.outputs[0], self.fgraph.inputs[0])
            self.fgraph.add_output(grad_wrt_x)

        if hess:
            hess_wrt_x = jacobian(
                self.fgraph.outputs[-1],
                self.fgraph.inputs[0],
                vectorize=use_vectorized_jac,
            )
            self.fgraph.add_output(hess_wrt_x)

        self.fgraph = self.fgraph.freeze()

        self.jac = jac
        self.hess = hess
        self.hessp = hessp

        self.method = method
        self.optimizer_kwargs = optimizer_kwargs if optimizer_kwargs is not None else {}
        self._fn = None
        self._fn_wrapped = None

    def connection_pattern(self, node=None):
        return _optimizer_connection_pattern(self.fgraph, is_minimization=True)

    def __str__(self):
        str_args = ", ".join(
            [
                f"{arg}={getattr(self, arg)}"
                for arg in ["method", "jac", "hess", "hessp"]
            ]
        )
        return f"{self.__class__.__name__}({str_args})"

    def perform(self, node, inputs, outputs):
        global optimize
        if optimize is None:
            import scipy.optimize as optimize

        f = self.fn_wrapped
        x0, *args = inputs

        res = optimize.minimize(
            fun=f.value_and_grad if self.jac else f.value,
            jac=self.jac,
            x0=x0,
            args=tuple(args),
            hess=f.hess if self.hess else None,
            method=self.method,
            **self.optimizer_kwargs,
        )

        f.clear_cache()

        outputs[0][0] = res.x.reshape(x0.shape).astype(x0.dtype)
        outputs[1][0] = np.bool_(res.success)

    def pullback(self, inputs, outputs, output_grads):
        x, *args = inputs
        x_star, _success = outputs
        output_grad, _ = output_grads

        d_xstar_d_theta = self.compute_implicit_gradients(
            x_star=x_star,
            args=args,
            output_grad=output_grad,
            is_minimization=True,
        )

        return [zeros_like(x), *d_xstar_d_theta]


def pack_inputs_of_objective(
    objective: TensorVariable,
    x: TensorVariable | Sequence[TensorVariable],
) -> tuple[TensorVariable, list[TensorVariable] | None, TensorVariable]:
    """
    Packs inputs `x` into a single tensor if `x` is a sequence of tensors, and rewrites the `objective` graph
    to use the packed input. Also returns the packed shapes for unpacking the output later.

    If `x` is a single tensor, it is returned as is, and no rewriting is done.
    """
    packed_shapes = None

    if not isinstance(x, Sequence):
        x = [x]

    if len(x) == 1 and x[0].ndim <= 1:
        packed_input = x[0]
    else:
        packed_input, packed_shapes = pack(*x)
        unpacked_output = unpack(packed_input, packed_shapes)

        objective = graph_replace(
            objective,
            {
                xi: ui.astype(xi.type.dtype)
                if not (isinstance(xi.type, ScalarType))
                else scalar_from_tensor(ui.astype(xi.type.dtype))
                for xi, ui in zip(x, unpacked_output)
            },
        )

    return packed_input, packed_shapes, objective


def minimize(
    objective: TensorVariable,
    x: TensorVariable | Sequence[TensorVariable],
    method: str = "BFGS",
    jac: bool = True,
    hess: bool = False,
    use_vectorized_jac: bool = False,
    optimizer_kwargs: dict | None = None,
) -> tuple[TensorVariable | tuple[TensorVariable, ...], TensorVariable]:
    """
    Minimize a scalar objective function using scipy.optimize.minimize.

    Parameters
    ----------
    objective : TensorVariable
        The objective function to minimize. This should be a pytensor variable representing a scalar value.
    x: TensorVariable or list of TensorVariable
        The variable or variables with respect to which the objective function is minimized. It must be an input to the
        computational graph of `objective`.
    method: str, optional
        The optimization method to use. Default is "BFGS". See scipy.optimize.minimize for other options.
    jac: bool, optional
        Whether to compute and use the gradient of the objective function with respect to x for optimization.
        Default is True.
    hess: bool, optional
        Whether to compute and use the Hessian of the objective function with respect to x for optimization.
        Default is False. Note that some methods require this, while others do not support it.
    use_vectorized_jac: bool, optional
        Whether to use a vectorized graph (vmap) to compute the jacobian (and/or hessian) matrix. If False, a
        scan will be used instead. This comes down to a memory/compute trade-off. Vectorized graphs can be faster,
        but use more memory. Default is False.
    optimizer_kwargs
        Additional keyword arguments to pass to scipy.optimize.minimize

    Returns
    -------
    solution: TensorVariable or tuple of TensorVariable
        The optimized value of each of inputs in `x` that minimizes `objective(x, *args)`. If the success flag
        is False, this will be the final state of the minimization routine, but not necessarily a minimum.

    success: TensorVariable
        Symbolic boolean flag indicating whether the minimization routine reported convergence to a minimum
        value, based on the requested convergence criteria.
    """
    objective = as_tensor_variable(objective)

    packed_input, packed_shapes, objective = pack_inputs_of_objective(objective, x)
    args = _find_optimization_parameters(objective, packed_input)
    objective, args, outer_args = _tensorize_scalar_parameters(objective, args)

    minimize_op = MinimizeOp(
        packed_input,
        *args,
        objective=objective,
        method=method,
        jac=jac,
        hess=hess,
        use_vectorized_jac=use_vectorized_jac,
        optimizer_kwargs=optimizer_kwargs,
    )

    solution, success = minimize_op(packed_input, *outer_args)

    if packed_shapes is not None:
        solution = unpack(solution, packed_shapes)

    return solution, success


class RootScalarOp(ScipyScalarWrapperOp):
    _scipy_props = ("method", "jac", "hess", "optimizer_kwargs")

    def __init__(
        self,
        variables: TensorVariable,
        *args: Variable,
        equation: TensorVariable,
        method: str,
        jac: bool = False,
        hess: bool = False,
        optimizer_kwargs=None,
    ):
        if not (isinstance(variables, TensorVariable) and variables.ndim == 0):
            raise ValueError(
                "The variable `x` must be a scalar (0-dimensional) tensor for root_scalar."
            )
        if not (isinstance(equation, TensorVariable) and equation.ndim == 0):
            raise ValueError(
                "The equation must be a scalar (0-dimensional) tensor for root_scalar."
            )
        if variables not in ancestors([equation]):
            raise ValueError(
                "The variable `variables` must be an input to the computational graph of the equation."
            )

        self.fgraph = FunctionGraph([variables, *args], [equation])

        if jac:
            f_prime = grad(self.fgraph.outputs[0], self.fgraph.inputs[0])
            self.fgraph.add_output(f_prime)

        if hess:
            if not jac:
                raise ValueError(
                    "Cannot set `hess=True` without `jac=True`. No methods use second derivatives without also"
                    " using first derivatives."
                )
            f_double_prime = grad(self.fgraph.outputs[-1], self.fgraph.inputs[0])
            self.fgraph.add_output(f_double_prime)

        self.fgraph = self.fgraph.freeze()

        self.method = method
        self.optimizer_kwargs = optimizer_kwargs if optimizer_kwargs is not None else {}
        self.jac = jac
        self.hess = hess

        self._fn = None
        self._fn_wrapped = None

    def connection_pattern(self, node=None):
        return _optimizer_connection_pattern(self.fgraph, is_minimization=False)

    def __str__(self):
        str_args = ", ".join(
            [f"{arg}={getattr(self, arg)}" for arg in ["method", "jac", "hess"]]
        )
        return f"{self.__class__.__name__}({str_args})"

    def perform(self, node, inputs, outputs):
        global optimize
        if optimize is None:
            import scipy.optimize as optimize

        f = self.fn_wrapped
        f.clear_cache()
        # f.copy_x = True

        variables, *args = inputs

        res = optimize.root_scalar(
            f=f.value,
            fprime=f.grad if self.jac else None,
            fprime2=f.hess if self.hess else None,
            x0=variables,
            args=tuple(args),
            method=self.method,
            **self.optimizer_kwargs,
        )

        outputs[0][0] = np.array(res.root)
        outputs[1][0] = np.bool_(res.converged)

    def pullback(self, inputs, outputs, output_grads):
        x, *args = inputs
        x_star, _ = outputs
        output_grad, _ = output_grads

        d_xstar_d_theta = self.compute_implicit_gradients(
            x_star=x_star,
            args=args,
            output_grad=output_grad,
            is_minimization=False,
        )

        return [zeros_like(x), *d_xstar_d_theta]


def root_scalar(
    equation: TensorVariable,
    variable: TensorVariable,
    method: str = "secant",
    jac: bool = False,
    hess: bool = False,
    optimizer_kwargs: dict | None = None,
) -> tuple[TensorVariable, TensorVariable]:
    """
    Find roots of a scalar equation using scipy.optimize.root_scalar.

    Parameters
    ----------
    equation : TensorVariable
        The equation for which to find roots. This should be a PyTensor variable representing a single equation in one
        variable. The function will find `variables` such that `equation(variables, *args) = 0`.
    variable : TensorVariable
        The variable with respect to which the equation is solved. It must be a scalar and an input to the
        computational graph of `equation`.
    method : str, optional
        The root-finding method to use. Default is "secant". See `scipy.optimize.root_scalar` for other options.
    jac : bool, optional
        Whether to compute and use the first derivative of the equation with respect to `variables`.
        Default is False. Some methods require this.
    hess : bool, optional
        Whether to compute and use the second derivative of the equation with respect to `variables`.
        Default is False. Some methods require this.
    optimizer_kwargs : dict, optional
        Additional keyword arguments to pass to `scipy.optimize.root_scalar`.

    Returns
    -------
    solution: TensorVariable
        The final state of the root-finding routine. When `success` is True, this is the value of `variables` that
        causes `equation` to evaluate to zero. Otherwise it is the final state returned by the root-finding
        routine, but not necessarily a root.

    success: TensorVariable
        Boolean indicating whether the root-finding was successful. If True, the solution is a root of the equation
    """
    args = _find_optimization_parameters(equation, variable)
    equation, args, outer_args = _tensorize_scalar_parameters(equation, args)

    root_scalar_op = RootScalarOp(
        variable,
        *args,
        equation=equation,
        method=method,
        jac=jac,
        hess=hess,
        optimizer_kwargs=optimizer_kwargs,
    )

    solution, success = root_scalar_op(variable, *outer_args)

    return solution, success


class RootOp(ScipyVectorWrapperOp):
    # eq/hash key on the (frozen) inner graph plus these props -- keying on the
    # inner graph is what keeps RootOps of different equations distinct (an
    # earlier ``__props__ = ("method", "jac")`` ignored it and merged them).
    _scipy_props = ("method", "jac", "use_vectorized_jac", "optimizer_kwargs")

    def __init__(
        self,
        variables: TensorVariable,
        *args: Variable,
        equations: TensorVariable,
        method: str,
        jac: bool = True,
        optimizer_kwargs: dict | None = None,
        use_vectorized_jac: bool = False,
    ):
        if not isinstance(variables, TensorVariable):
            raise ValueError("The variable `variables` must be a tensor for root.")
        if not isinstance(equations, TensorVariable):
            raise ValueError("The equations must be a tensor for root.")
        if variables.ndim != equations.ndim:
            raise ValueError(
                "The variable `variables` must have the same number of dimensions as the equations."
            )
        if variables not in ancestors([equations]):
            raise ValueError(
                "The variable `variables` must be an input to the computational graph of the equations."
            )

        self.fgraph = FunctionGraph([variables, *args], [equations])
        self.use_vectorized_jac = use_vectorized_jac

        if jac:
            jac_wrt_x = jacobian(
                self.fgraph.outputs[0],
                self.fgraph.inputs[0],
                vectorize=use_vectorized_jac,
            )
            self.fgraph.add_output(atleast_2d(jac_wrt_x))

        self.fgraph = self.fgraph.freeze()

        self.jac = jac

        self.method = method
        self.optimizer_kwargs = optimizer_kwargs if optimizer_kwargs is not None else {}
        self._fn = None
        self._fn_wrapped = None

    def connection_pattern(self, node=None):
        return _optimizer_connection_pattern(self.fgraph, is_minimization=False)

    def __str__(self):
        str_args = ", ".join(
            [f"{arg}={getattr(self, arg)}" for arg in ["method", "jac"]]
        )
        return f"{self.__class__.__name__}({str_args})"

    def compile_fn(self, mode):
        fgraph = self.fgraph.unfreeze()
        variables, *args = fgraph.inputs
        outputs = fgraph.outputs

        if variables.ndim > 0:
            new_root_variables = variables
            new_outputs = outputs
        else:
            # If the user passes a scalar optimization problem to root, scipy will automatically upcast it to
            # a 1d array. The inner function needs to be adjusted to handle this.
            new_root_variables = tensor(
                name="variables_vector", shape=(1,), dtype=variables.type.dtype
            )
            new_variables = new_root_variables.squeeze()

            new_outputs = graph_replace(outputs, {variables: new_variables})

        return self.link_fgraph(
            FunctionGraph([new_root_variables, *args], new_outputs, clone=False), mode
        )

    def perform(self, node, inputs, outputs):
        global optimize
        if optimize is None:
            import scipy.optimize as optimize

        f = self.fn_wrapped
        f.clear_cache()
        f.copy_x = True

        variables, *args = inputs

        res = optimize.root(
            fun=f,
            jac=self.jac,
            x0=variables,
            args=tuple(args),
            method=self.method,
            **self.optimizer_kwargs,
        )

        # There's a reshape here to cover the case where variables is a scalar. Scipy will still return a
        # (1, 1) matrix in in this case, which causes errors downstream (since pytensor expects a scalar).
        outputs[0][0] = res.x.reshape(variables.shape).astype(variables.dtype)
        outputs[1][0] = np.bool_(res.success)

    def pullback(self, inputs, outputs, output_grads):
        x, *args = inputs
        x_star, _ = outputs
        output_grad, _ = output_grads

        d_xstar_d_theta = self.compute_implicit_gradients(
            x_star=x_star,
            args=args,
            output_grad=output_grad,
            is_minimization=False,
        )

        return [zeros_like(x), *d_xstar_d_theta]


def root(
    equations: TensorVariable,
    variables: TensorVariable | Sequence[TensorVariable],
    method: str = "hybr",
    jac: bool = True,
    use_vectorized_jac: bool = False,
    optimizer_kwargs: dict | None = None,
) -> tuple[TensorVariable | Sequence[TensorVariable], TensorVariable]:
    """
    Find roots of a system of equations using scipy.optimize.root.

    Parameters
    ----------
    equations : TensorVariable
        The system of equations for which to find roots. This should be a PyTensor variable representing a
        vector (or scalar) value. The function will find `variables` such that `equations(variables, *args) = 0`.
    variables : TensorVariable
        The variable(s) with respect to which the system of equations is solved. It must be an input to the
        computational graph of `equations` and have the same number of dimensions as `equations`.
    method : str, optional
        The root-finding method to use. Default is "hybr". See `scipy.optimize.root` for other options.
    jac : bool, optional
        Whether to compute and use the Jacobian of the `equations` with respect to `variables`.
        Default is True. Most methods require this.
    use_vectorized_jac: bool, optional
        Whether to use a vectorized graph (vmap) to compute the jacobian matrix. If False, a scan will be used instead.
        This comes down to a memory/compute trade-off. Vectorized graphs can be faster, but use more memory.
        Default is False.
    optimizer_kwargs : dict, optional
        Additional keyword arguments to pass to `scipy.optimize.root`.

    Returns
    -------
    solution: TensorVariable
        The final state of the root-finding routine. When `success` is True, this is the value of `variables` that
        causes all `equations` to evaluate to zero. Otherwise it is the final state returned by the root-finding
        routine, but not necessarily a root.

    success: TensorVariable
        Boolean indicating whether the root-finding was successful. If True, the solution is a root of the equation
    """
    packed_variables, packed_shapes, equations = pack_inputs_of_objective(
        equations, variables
    )
    args = _find_optimization_parameters(equations, packed_variables)
    equations, args, outer_args = _tensorize_scalar_parameters(equations, args)

    root_op = RootOp(
        packed_variables,
        *args,
        equations=equations,
        method=method,
        jac=jac,
        optimizer_kwargs=optimizer_kwargs,
        use_vectorized_jac=use_vectorized_jac,
    )

    solution, success = root_op(packed_variables, *outer_args)
    if packed_shapes is not None:
        solution = unpack(solution, packed_shapes)

    return solution, success


__all__ = ["minimize", "minimize_scalar", "root", "root_scalar"]
