"""The base models used to specify the TDAstro computation graph.

The computation graph is composed of ParameterizedNodes which compute random variables
(called "model parameters" or "parameters" for short) and encode the dependencies between
them. Model parameters are different from object variables in that they are *not* stored
within the object, but rather programmatically set in an external graph_state dictionary by
sampling the graph. The user does not need to access the internals of graph_state directly,
but rather can treat it as an opaque object to pass around. The dictionary can hold either
individual values (often floats) or arrays of samples.

All model parameters (random variables in the probabilistic graph) must be added using
using the ParameterizedNode.add_parameter() function. This allows the graph to track
which parameters to update and how to set them. Parameter's values can be set from a few
sources:
1) A constant (Example: A given standard deviation for a noise model)
2) A static function or method (which does not have variables that are resampled).
3) The result of evaluating a FunctionNode, which provides a computation using other
   parameters in the graph.
4) The parameters of another ParameterizedNode.

We say that parameter X is dependent on parameter Y if the value of Y is necessary
to compute the value of X. For example if X is set by evaluating a FunctionNode that
uses parameter Y in the computation, X is dependent on Y. The dependencies impose an
ordering of model parameters in the graph. Y must be computed before X.

ParameterNodes provide semantic groupings of individual parameters. For example we may
have a ParameterNode representing the information needed for a Type Ia supernova.
That node's parameters would include the variables needed to evaluate the supernova's
lightcurve. Each of these parameters might depend on parameters in other nodes, such
as those of the host galaxy.

The execution graph is processed by starting at the final node, examining each
model parameter for that node, and recursively proceeding 'up' the graph for any
of its parameters that has a dependency. For example the function.

f(a, b) = x
g(c) = y
h(x, y) = z

would form the graph:

a -\
    x - \
b -/     \
          z
c -- y -- /

where z is the 'bottom' node. Parameters a, b, and c would be at the 'top' of the
graph because they have no dependencies.  Such parameters are set by constants or
static functions.
"""

from functools import partial

import numpy as np

from tdastro.graph_state import GraphState


class ParameterSource:
    """ParameterSource specifies the information about where a ParameterizedNode should
    get the value for a given parameter.

    Attributes
    ----------
    parameter_name : str
        The name of the parameter within the node (short name).
    node_name : str
        The name of the parent node.
    source_type : int
        The type of source as defined by the class variables.
        Default = 0
    value : any
        The information that actually sets the parameter. Either a constant
        or the attribute name of a dependency node.
    dependency : ParameterizedNode or None
        The node on which this parameter is dependent
    allow_gradient : bool
        Allow gradients to be computed at this variable.
        Default = False
    """

    # Class variables for the source enum.
    UNDEFINED = 0
    CONSTANT = 1
    MODEL_PARAMETER = 2
    FUNCTION_NODE = 3
    COMPUTE_OUTPUT = 4

    def __init__(self, parameter_name, source_type=0, node_name=""):
        self.parameter_name = parameter_name
        self.node_name = node_name
        self.source_type = source_type
        self.allow_gradient = False
        self.value = None
        self.dependency = None

    def set_as_constant(self, value, allow_gradient=True):
        """Set the parameter as a constant value.

        Parameters
        ----------
        value : any
            The constant value to use.
        allow_gradient : bool
            Allow a gradient to be computed at this variable.
            Default = True
        """
        if callable(value):
            raise ValueError(f"Using set_as_constant on callable {value}")

        self.source_type = ParameterSource.CONSTANT
        self.allow_gradient = allow_gradient
        self.dependency = None
        self.value = value

    def set_as_parameter(self, dependency, param_name):
        """Set the parameter as a model parameter of another node.  This is
        used for chaining, such as when an object's ra depends on its host's ra.

        Parameters
        ----------
        dependency : ParameterizedNode
            The node in which to access the attribute.
        param_name : str
            The name of the parameter to access.
        """
        self.source_type = ParameterSource.MODEL_PARAMETER
        self.allow_gradient = False
        self.dependency = dependency
        self.value = param_name

    def set_as_function(self, dependency, param_name="function_node_result"):
        """Set the parameter as the result of a FunctionNode (that is not
        the current node).

        Parameters
        ----------
        dependency : ParameterizedNode
            The node in which to access the attribute.
        param_name : str
            The name of where the result is stored in the FunctionNode.
        """
        self.source_type = ParameterSource.FUNCTION_NODE
        self.allow_gradient = False
        self.dependency = dependency
        self.value = param_name

    def set_as_compute_output(self, param_name="function_node_result"):
        """Set the parameter the output of this current node's compute() method.

        This needs to be separate from FUNCTION_NODE type (set_as_function) because
        the sampling function needs to know to call the current node's compute()
        method after the other parameters have been sampled.

        Parameters
        ----------
        dependency : ParameterizedNode
            The node in which to access the attribute.
        param_name : str
            The name of where the result is stored in the FunctionNode.
        """
        self.source_type = ParameterSource.COMPUTE_OUTPUT
        self.allow_gradient = False
        self.value = param_name


class ParameterizedNode:
    """Any model that uses parameters that can be set by constants,
    functions, or other parameterized nodes.

    ParameterizedNodes do not store values directly, but rather provide a recipe
    for how to generate the parameters' values.  The sampled values are read from
    and written to a GraphState object that stores all the parameter values for
    all the nodes.

    Attributes
    ----------
    node_label : str
        An optional human readable identifier (name) for the current node.
    node_string : str
        The full string used to identify a node. This is a combination of the nodes position
        in the graph (if known), node_label (if provided), and class information. This is
        used to access the parameters for this node in the graph_state.
    setters : dict
        A dictionary mapping the parameters' names to information about the setters
        (ParameterSource). The model parameters are stored in the order in which they
        need to be set.
    node_pos : int or None
        A unique ID number for each node in the graph indicating its position.
        Assigned during resampling or set_graph_positions(). This is required to resolve
        naming collisions so we do not overwrite parameters from other nodes.

    Parameters
    ----------
    node_label : str, optional
        An identifier (or name) for the current node.
    **kwargs : dict, optional
        Any additional keyword arguments.
    """

    def __init__(self, node_label=None, **kwargs):
        self.setters = {}
        self.node_label = node_label
        self.node_pos = None
        self.node_string = None

        # Give the node a temporary name.
        self._update_node_string()

    def __str__(self):
        """Return the string representation of the node."""
        return self.node_string

    def _update_node_string(self, new_str=None):
        """Update the node's string.

        Parameters
        ----------
        new_str : str, optional
            The new node string. If not provided, the node_string
            is automatically computed from the other node information.
        """
        if self.node_label is not None:
            # If a label is given, just use that. It overrides even the new_str.
            self.node_string = self.node_label
        elif new_str is not None:
            self.node_string = new_str
        else:
            # Otherwise use a combination of the node's class and position.
            pos_string = f"_{self.node_pos}" if self.node_pos is not None else ""
            self.node_string = f"{self.__class__.__name__}{pos_string}"

        # Update the node_name of all node's parameter setters.
        for _, setter_info in self.setters.items():
            setter_info.node_name = self.node_string

    def set_graph_positions(self, seen_nodes=None):
        """Force an update of the graph structure (numbering of each node).

        Parameters
        ----------
        seen_nodes : set, optional
            A set of nodes that have already been processed to prevent infinite loops.
            Caller should not set.
        """
        # Make sure that we do not process the same nodes multiple times.
        if seen_nodes is None:
            seen_nodes = set()
        if self in seen_nodes:
            return
        seen_nodes.add(self)

        # Update the node's position in the graph and its string.
        self.node_pos = len(seen_nodes) - 1
        self._update_node_string()

        # Recursively update any direct dependencies.
        for setter_info in self.setters.values():
            if setter_info.dependency is not None and setter_info.dependency is not self:
                setter_info.dependency.set_graph_positions(seen_nodes)

    def list_params(self):
        """Return a list of this node's parameterized values.

        Returns
        -------
        names : list[str]
            The name of all of the parameterized values for this node.
        """
        return list(self.setters.keys())

    def has_valid_param(self, name):
        """Check whether the node has a given parameterized value and that it is not
        always set to None.

        Parameters
        ----------
        name : str
            The name of the parameter.

        Returns
        -------
        contains : bool
            Whether the node contains a given parameter and it is not always None.
        """
        if name not in self.setters:
            return False

        setter = self.setters[name]
        return not (setter.source_type == ParameterSource.CONSTANT and setter.value is None)

    def get_param(self, graph_state, name, default=None):
        """Get the value of a parameter stored in this node or a default value.

        Note
        ----
        This is an optional helper function that accesses the internals of graph_state
        using the node's information (e.g. its hash value).

        Parameters
        ----------
        graph_state : dict
            The dictionary of graph state information.
        name : str
            The parameter name to query.
        default : any
            The default value to return if the parameter is not in GraphState.

        Returns
        -------
        any
            The parameter value or the default.

        Raises
        ------
        ValueError if graph_state is None.
        """
        if graph_state is None:
            raise ValueError(f"Unable to look up parameter={name}. No graph_state given.")
        if self.node_string in graph_state and name in graph_state[self.node_string]:
            return graph_state[self.node_string][name]
        return default

    def get_local_params(self, graph_state):
        """Get a dictionary of all parameters local to this node.

        Note
        ----
        This is an optional helper function that accesses the internals of graph_state
        using the node's information (e.g. its hash value).

        Parameters
        ----------
        graph_state : GraphState
            An object mapping graph parameters to their values.

        Returns
        -------
        result : dict
            A dictionary mapping the parameter name to its value.

        Raises
        ------
        KeyError if no parameters have been set for this node.
        ValueError if graph_state is None.
        """
        if graph_state is None:
            raise ValueError("No graph_state given.")
        return graph_state[self.node_string]

    def set_parameter(self, name, value=None, **kwargs):
        """Set the source of a single *existing* parameter in the ParameterizedNode.

        Parameters within a node are actually a mapping of the parameter name
        to a ParameterSource object that indicates how they are set during sampling.
        Some parameters may be set as a constant value while others may be set by
        evaluating a function that depends on other parameters.

        Notes
        -----
        * Does NOT set an initial value for the model parameter. The user must
          sample the parameters for this to be set.
        * The model parameters are stored in the order in which they are added.

        Parameters
        ----------
        name : str
            The parameter name to add.
        value : any, optional
            The information to use to set the parameter. Can be a constant,
            function, ParameterizedNode, or self.
        **kwargs : dict, optional
           All other keyword arguments, possibly including the parameter setters.

        Raises
        ------
        Raise a KeyError if there is a parameter collision or the parameter
        cannot be found.
        Raise a ValueError if the setter type is not supported.
        """
        # Set the node's position in the graph to None to indicate that the
        # structure might have changed. It needs to be updated with set_graph_positions().
        self.node_pos = None

        # Check for parameter has been added and if so, find the index. All parameters must
        # be added first with "add_parameter()".
        if name not in self.setters:
            raise KeyError(
                f"Tried to set parameter '{name}' that has not been added to node {self.node_string}."
            ) from None

        if value is None and name in kwargs:
            # The value wasn't set, but the name is in kwargs.
            value = kwargs[name]

        if callable(value):
            if "__self__" in value.__dir__() and isinstance(value.__self__, ParameterizedNode):
                # Case 1a: This is a method attached to another ParameterizedNode.
                # Check if this is a getter method, including one that might have been automatically
                # created in add_parameter(). If this is a getter, we access the parameter's sampled
                # values directly to save a function call.
                method_name = value.__name__
                parent = value.__self__
                if method_name in parent.setters:
                    self.setters[name].set_as_parameter(parent, method_name)
                else:
                    raise ValueError(
                        f"Trying to set parameter '{name}' to the {type(parent)}.{method_name}, "
                        f"but unable to find that parameter in class {type(parent)}."
                    )
            else:
                # Case 1b: This is a general function or callable method from another object.
                # We treat it as static (we don't resample the other object) and
                # wrap it in a FunctionNode.
                func_node = FunctionNode(value, **kwargs)
                self.setters[name].set_as_function(func_node)
        elif isinstance(value, FunctionNode):
            # Case 2: We are using the result of a computation of the function node.
            # If the FunctionNode has names outputs that match the variable, use that.
            output_name = name if name in value.outputs else "function_node_result"
            self.setters[name].set_as_function(value, output_name)
        elif isinstance(value, ParameterizedNode):
            # Case 3: We are trying to access a parameter of a ParameterizedNode
            # with the same name.
            if value == self:
                raise ValueError(f"Parameter '{name}' is recursively assigned to self.{name}.")
            if name not in value.setters:
                raise ValueError(f"Parameter '{name}' missing from {str(value)}.")
            self.setters[name].set_as_parameter(value, name)
        else:
            # Case 4: The value is constant (including None).
            self.setters[name].set_as_constant(value)

    def set_allow_gradient(self, name, allow_gradient):
        """Turn on or off the ability to compute a gradient for this variable.

        Parameters
        ----------
        name : str
            The parameter name to modify.
        allow_gradient : bool
            The new setting for allow_gradient.
        """
        self.setters[name].allow_gradient = allow_gradient

    def add_parameter(self, name, value=None, allow_gradient=None, **kwargs):
        """Add a single *new* parameter to the ParameterizedNode.

        Notes
        -----
        * Checks multiple sources in the following order: Manually specified value,
          an entry in kwargs, or None.
        * Does NOT set an initial value for the model parameter. The user must
          sample the parameters for this to be set.
        * The model parameters are stored in the order in which they are added.

        Parameters
        ----------
        name : str
            The parameter name to add.
        value : any, optional
            The information to use to set the parameter. Can be a constant,
            function, ParameterizedNode, or self.
        allow_gradient : bool or None
            Allow gradients to be computed for this variable. If set to None uses the default
            for the setter type (True for constant and False for everything else).
            Default = None
        **kwargs : dict, optional
           All other keyword arguments, possibly including the parameter setters.

        Raises
        ------
        Raise a KeyError if there is a parameter collision or the parameter
        cannot be found.
        """
        # Check for parameter collision and add a place holder value to the 'setters' dictionary.
        if hasattr(self, name) and name not in self.setters:
            raise KeyError(
                f"Parameter name '{name}' conflicts with a predefined model parameter "
                f"or class attribute in {self.node_string}"
            )
        if self.setters.get(name, None) is not None:
            raise KeyError(f"Duplicate parameter set: '{name}' in {self.node_string}")

        # Add an entry for the setter function and fill in the remaining information using
        # set_parameter(). We add an initial (dummy) value here to indicate that this parameter
        # exists and was added via add_parameter().
        self.setters[name] = ParameterSource(
            parameter_name=name,
            source_type=ParameterSource.UNDEFINED,
            node_name=str(self),
        )
        self.set_parameter(name, value, **kwargs)

        # Check if we should override allow_gradient.
        if allow_gradient is not None:
            self.setters[name].allow_gradient = allow_gradient

        # Create a callable getter function with the same name as the parameter.
        # This function allows us to reference the parameter as object.parameter_name
        # for chaining without copying the value. For example, if my_node_1, is a
        # ParameterizedNode with a parameter x, we can do:
        #   my_node_2 = ParameterizedNode(y=my_node_1.x)
        # and my_node_2 will know to use the sampled values of x from my_node_1
        # (as opposed to the setter for x).
        #
        # We override the __self__ and __name__ attributes so it looks like method of
        # this object and the assignment y=my_node_1.x doesn't do a copy of the value.
        def getter(graph_state):
            return graph_state[getter.__self__.node_string][getter.__name__]

        getter.__self__ = self
        getter.__name__ = name
        setattr(self, name, getter)

    def compute(self, graph_state, rng_info=None, **kwargs):
        """Placeholder for a general compute function, which is called at the end
        of the sampling process and can produce derived parameters. This function
        is the main processing step in a FunctionNode.

        Parameters
        ----------
        graph_state : GraphState
            An object mapping graph parameters to their values. This object is modified
            in place as it is sampled.
        rng_info : numpy.random._generator.Generator, optional
            A given numpy random number generator to use for this computation. If not
            provided, the function uses the node's random number generator.
        **kwargs : dict, optional
            Additional function arguments.
        """
        return None

    def _sample_helper(self, graph_state, seen_nodes, rng_info=None):
        """Internal recursive function to sample the model's underlying parameters
        if they are provided by a function or ParameterizedNode. All sampled
        parameters for all nodes are stored in the graph_state dictionary, which is
        modified in-place.

        Parameters
        ----------
        graph_state : GraphState
            An object mapping graph parameters to their values. This object is modified
            in place as it is sampled.
        seen_nodes : dict
            A dictionary mapping nodes strings seen during this sampling run to their object.
            Used to avoid sampling nodes multiple times and to validity check the graph.
        rng_info : numpy.random._generator.Generator, optional
            A given numpy random number generator to use for this computation. If not
            provided, the function uses the node's random number generator.

        Raises
        ------
        Raise a KeyError if the sampling encounters an error with the order of dependencies.
        """
        node_str = str(self)
        if node_str in seen_nodes:
            if seen_nodes[node_str] != self:
                raise ValueError(
                    f"Duplicate node label '{node_str}'. Every node must have a unique label. "
                    "This most often happens when the node_label parameter is set directly."
                )
            return  # Nothing to do
        seen_nodes[node_str] = self

        # Run through each parameter and sample it based on the given recipe.
        # As of Python 3.7 dictionaries are guaranteed to preserve insertion ordering,
        # so this will iterate through model parameters in the order they were inserted.
        any_compute = False
        for name, setter in self.setters.items():
            # Check if we need to sample this parameter's dependency node.
            if setter.dependency is not None and setter.dependency != self:
                setter.dependency._sample_helper(graph_state, seen_nodes, rng_info=rng_info)

            # Set the result from the correct source.
            if setter.source_type == ParameterSource.CONSTANT:
                if graph_state.num_samples == 1:
                    graph_state.set(self.node_string, name, setter.value)
                else:
                    repeated_value = np.array([setter.value] * graph_state.num_samples)
                    graph_state.set(self.node_string, name, repeated_value)
            elif setter.source_type == ParameterSource.MODEL_PARAMETER:
                graph_state.set(
                    self.node_string,
                    name,
                    graph_state[setter.dependency.node_string][setter.value],
                )
            elif setter.source_type == ParameterSource.FUNCTION_NODE:
                graph_state.set(
                    self.node_string,
                    name,
                    graph_state[setter.dependency.node_string][setter.value],
                )
            elif setter.source_type == ParameterSource.COMPUTE_OUTPUT:
                # Computed parameters are set only after all the other (input) parameters.
                any_compute = True
            else:
                raise ValueError(f"Invalid ParameterSource type {setter.source_type}")

        # If this is a function node and the parameters depend on the result of its own computation
        # call the compute function to fill them in.
        if any_compute:
            self.compute(graph_state, rng_info)

    def sample_parameters(self, given_args=None, num_samples=1, rng_info=None):
        """Sample the model's underlying parameters if they are provided by a function
        or ParameterizedNode.

        Parameters
        ----------
        given_args : dict, optional
            A dictionary representing the given arguments for this sample run.
            This can be used as the JAX PyTree for differentiation.
        num_samples : int
            A count of the number of samples to compute.
            Default: 1
        rng_info : numpy.random._generator.Generator, optional
            A given numpy random number generator to use for this computation. If not
            provided, the function uses the node's random number generator.

        Returns
        -------
        graph_state : GraphState
            A dictionary of dictionaries mapping node->hash, variable_name to either a
            value or array of values. This data structure is modified in place to represent
            the model's state(s).

        Raises
        ------
        Raise a ValueError the sampling encounters a problem with the order of dependencies.
        """
        # If the graph structure has never been set, do that now.
        if self.node_pos is None:
            nodes = set()
            self.set_graph_positions(seen_nodes=nodes)

        # Create space for the results and set all the given_args as fixed parameters.
        results = GraphState(num_samples)
        if given_args is not None:
            results.update(given_args, all_fixed=True)

        # Resample the nodes. All information is stored in the returned results dictionary.
        seen_nodes = {}
        self._sample_helper(results, seen_nodes, rng_info=rng_info)
        return results

    def build_pytree(self, graph_state, partial=None):
        """Build a JAX PyTree representation of the variables in this graph.

        Parameters
        ----------
        graph_state : dict
            A dictionary of dictionaries mapping node->hash, variable_name to value.
            This data structure is modified in place to represent the current state.
        partial : dict
            The partial results so far. This is modified in place by the function.
            A dictionary mapping node name to a dictionary mapping each variable's name
            to its value.
            Default : None

        Returns
        -------
        values : dict
            A dictionary mapping node name to a dictionary mapping each variable's name
            to its value.
        """
        # Check if the node might have incomplete information.
        if self.node_pos is None:
            raise ValueError(
                f"Node {self.node_string} is missing position. You must call "
                "set_graph_positions() before building a pytree."
            )

        # Skip nodes that we have already seen.
        if partial is None:
            partial = {}
        if self.node_string in partial:
            return partial

        # Add new values to the pytree, recursively exploring dependencies.
        partial[self.node_string] = {}
        for name, setter_info in self.setters.items():
            if setter_info.allow_gradient:
                # Anything wth allow_gradient == True goes in the PyTree.
                partial[self.node_string][name] = graph_state[self.node_string][name]
            elif setter_info.dependency is not None:
                # We only recursively check parameters above non-gradient nodes.
                partial = setter_info.dependency.build_pytree(graph_state, partial)
        return partial


class FunctionNode(ParameterizedNode):
    """A class to wrap functions and their argument settings.

    The node can compute the result using a given function (the func
    parameter) or through the compute() method. If no func=None
    then the user must override compute().

    Attributes
    ----------
    func : function or method or partial
        The function to call during an evaluation. If this is None
        you must override the compute() method directly.
    args_names : list
        A list of argument names to pass to the function.
    outputs : list of str
        The output model parameters of this function.

    Parameters
    ----------
    func : function or method
        The function to call during an evaluation.
    node_label : str, optional
        An identifier (or name) for the current node.
    outputs : list of str, optional
        The output model parameters of this function. If None uses
        a single model parameter result.
    fixed_params : dict, optional
        A dictionary mapping a parameter name in the function to its fixed value.
    **kwargs : dict, optional
        Any additional keyword arguments.

    Examples
    --------
    my_func = TDFunc(random.randint, a=1, b=10)
    value1 = my_func()      # Sample from default range
    value2 = my_func(b=20)  # Sample from extended range

    Note
    ----
    All the function's parameters that will be used need to be specified
    in either the default_args dict, object_args list, or as a kwarg in the
    constructor. Arguments cannot be first given during function call.
    For example the following will fail (because b is not defined in the
    constructor):

    my_func = TDFunc(random.randint, a=1)
    value1 = my_func(b=10.0)
    """

    def __init__(self, func, node_label=None, outputs=None, fixed_params=None, **kwargs):
        # We set the function before calling the parent class so we can use
        # the function's name (if needed).
        if fixed_params is not None and len(fixed_params) > 0:
            # Create a partial function with some of the parameters fixed.
            self.func = partial(func, **fixed_params)

            # We need to set the __name__ parameter because it is not preserved by partial.
            self.func.__name__ = func.__name__
        else:
            # Use the function as-is.
            self.func = func

        super().__init__(node_label=node_label, **kwargs)

        # Add all of the parameters from default_args or the kwargs.
        self.arg_names = []
        for key, value in kwargs.items():
            self.arg_names.append(key)
            self.add_parameter(key, value)

        # Add the output arguments.
        if not outputs:
            outputs = ["function_node_result"]
        self.outputs = outputs
        for name in outputs:
            # For output parameters we add a placeholder of None to set up the basic data, such as
            # the getter function and the entry in parameters. Then we change the
            # type to point to own result.
            self.add_parameter(name, None)
            self.setters[name].set_as_compute_output(param_name=name)

    def _non_func(self):
        """This function does nothing. This is used for FunctionNodes where the actual computation
        happens in an overloaded compute() function."""
        pass

    def _update_node_string(self, new_str=None):
        """Update the node's string. A Function node's string includes
        the function name in addition to the class name.
        """
        if new_str is None:
            pos_string = f"_{self.node_pos}" if self.node_pos is not None else ""
            fn_str = f":{self.func.__name__}" if self.func is not None else ""
            new_str = f"{self.__class__.__name__}{fn_str}{pos_string}"
        super()._update_node_string(new_str)

    def _build_inputs(self, graph_state, **kwargs):
        """Build the input arguments for the function.

        Parameters
        ----------
        graph_state : GraphState
            An object mapping graph parameters to their values. This object is modified
            in place as it is sampled.
        **kwargs : dict, optional
            Additional function arguments.

        Returns
        -------
        args : dict
            A dictionary of input argument to value.
        """
        args = {}
        for key in self.arg_names:
            if key in kwargs:
                args[key] = kwargs[key]
            else:
                args[key] = graph_state[self.node_string][key]
        return args

    def _save_results(self, results, graph_state):
        """Save the results to the graph state.

        Parameters
        ----------
        results : iterable
            The function's results.
        graph_state : GraphState
            An object mapping graph parameters to their values. This object is modified
            in place as it is sampled.
        """
        if len(self.outputs) == 1:
            graph_state.set(self.node_string, self.outputs[0], results)
        else:
            if len(results) != len(self.outputs):
                raise ValueError(
                    f"Incorrect number of results returned by {self.func.__name__}. "
                    f"Expected {len(self.outputs)}, but got {results}."
                )
            for i in range(len(self.outputs)):
                graph_state.set(self.node_string, self.outputs[i], results[i])

    def compute(self, graph_state, rng_info=None, **kwargs):
        """Execute the wrapped function.

        The input arguments are taken from the current graph_state and the outputs
        are written to graph_state.

        Parameters
        ----------
        graph_state : GraphState
            An object mapping graph parameters to their values. This object is modified
            in place as it is sampled.
        rng_info : numpy.random._generator.Generator, optional
            A given numpy random number generator to use for this computation. If not
            provided, the function uses the node's random number generator.
        **kwargs : dict, optional
            Additional function arguments.

        Returns
        -------
        results : any
            The result of the computation. This return value is provided so that testing
            functions can easily access the results.

        Raises
        ------
        ValueError is func attribute is None.
        """
        if self.func is None:
            raise ValueError(
                f"The FunctionNode {self.node_string}'s 'func' parameter is None. "
                "You need to either set func or override compute()."
            )

        # Build a dictionary of arguments for the function, call the function, and save
        # the results in the graph state.
        args = self._build_inputs(graph_state, **kwargs)
        results = self.func(**args)
        self._save_results(results, graph_state)
        return results

    def generate(self, given_args=None, num_samples=1, rng_info=None, **kwargs):
        """A helper function that regenerates the parameters for this nodes and the
        ones above it, then returns the the output or this individual node.

        This is used both for testing and for computing JAX gradients.

        Parameters
        ----------
        given_args : dict, optional
            A dictionary representing the given arguments for this sample run.
            This can be used as the JAX PyTree for differentiation.
        num_samples : int
            A count of the number of samples to compute.
            Default: 1
        rng_info : numpy.random._generator.Generator, optional
            A given numpy random number generator to use for this computation. If not
            provided, the function uses the node's random number generator.
        **kwargs : dict, optional
            Additional function arguments.
        """
        state = self.sample_parameters(given_args, num_samples, rng_info)

        # Get the result(s) of compute from the state object.
        if len(self.outputs) == 1:
            return self.get_param(state, self.outputs[0])

        results = []
        for output_name in self.outputs:
            results.append(self.get_param(state, output_name))
        return results
