r"""
    Allow an SGD to accept nested sequences of Parameters to optimize.

    This decorator can intepret the :code:`Parameter` objects in `btypes.py`,
    and can accept nested sequences of *any* structure of these objects to
    optimise!

    It can also optionally evaluate *random starts* (i.e. random starting
    candidates) if the parameter objects have been initialised with
    distributions. For this, an additional parameter is exposed in the
    :code:`minimizer` interface.

    Parameters
    ----------
    fun : callable
        objective function that takes in arbitrary ndarrays, floats or nested
        sequences of these.
    parameters : (nested) sequences of Parameter objects
        Initial guess of the parameters of the objective function
    nstarts : int, optional
        The number random starting candidates for optimisation to evaluate.
        This will only happen for :code:`nstarts > 0` and if at least one
        :code:`Parameter` object is random.

    Examples
    --------
    >>> from ..optimize import sgd
    >>> from ..btypes import Parameter, Bound

    Define a cost function that returns a pair. The first element is the cost
    value and the second element is the gradient represented by a sequence.
    Even if the cost is a function of a single variable, the gradient must be a
    sequence containing one element.

    >>> def cost(w, lambda_, data):
    ...     N = len(data)
    ...     y, X = data[:, 0], data[:, 1:]
    ...     y_est = X.dot(w)
    ...     ww = w.T.dot(w)
    ...     obj = (y - y_est).sum() / N + lambda_ * ww
    ...     gradw = - 2 * X.T.dot(y - y_est) / N + 2 * lambda_ * w
    ...     gradl = ww
    ...     return obj, [gradw, gradl]

    Augment the SGD optimizer to take structured inputs

    >>> new_sgd = structured_sgd(sgd)

    Data

    >>> y = np.linspace(1, 10, 100) + np.random.randn(100) + 1
    >>> X = np.array([np.ones(100), np.linspace(1, 100, 100)]).T
    >>> data = np.hstack((y[:, np.newaxis], X))

    Constant Initial values

    >>> w_0 = Parameter(np.array([1., 1.]), Bound())
    >>> lambda_0 = Parameter(.25, Bound())

    >>> res = new_sgd(cost, [w_0, lambda_0], data, batch_size=10,
    ...               eval_obj=True)
    >>> res_w, res_lambda = res.x

    Random Initial values

    >>> from scipy.stats import norm, gamma
    >>> w_0 = Parameter(norm(), Bound(), shape=(2,))
    >>> lambda_0 = Parameter(gamma(1.), Bound())

    >>> res = new_sgd(cost, [w_0, lambda_0], data, batch_size=10,
    ...               eval_obj=True, nstarts=100)
    >>> res_w, res_lambda = res.x

def structured_sgd(sgd):
    r"""
    Allow an SGD to accept nested sequences of Parameters to optimize.

    This decorator can intepret the :code:`Parameter` objects in `btypes.py`,
    and can accept nested sequences of *any* structure of these objects to
    optimise!

    It can also optionally evaluate *random starts* (i.e. random starting
    candidates) if the parameter objects have been initialised with
    distributions. For this, an additional parameter is exposed in the
    :code:`minimizer` interface.

    Parameters
    ----------
    fun : callable
        objective function that takes in arbitrary ndarrays, floats or nested
        sequences of these.
    parameters : (nested) sequences of Parameter objects
        Initial guess of the parameters of the objective function
    nstarts : int, optional
        The number random starting candidates for optimisation to evaluate.
        This will only happen for :code:`nstarts > 0` and if at least one
        :code:`Parameter` object is random.

    Examples
    --------
    >>> from ..optimize import sgd
    >>> from ..btypes import Parameter, Bound

    Define a cost function that returns a pair. The first element is the cost
    value and the second element is the gradient represented by a sequence.
    Even if the cost is a function of a single variable, the gradient must be a
    sequence containing one element.

    >>> def cost(w, lambda_, data):
    ...     N = len(data)
    ...     y, X = data[:, 0], data[:, 1:]
    ...     y_est = X.dot(w)
    ...     ww = w.T.dot(w)
    ...     obj = (y - y_est).sum() / N + lambda_ * ww
    ...     gradw = - 2 * X.T.dot(y - y_est) / N + 2 * lambda_ * w
    ...     gradl = ww
    ...     return obj, [gradw, gradl]

    Augment the SGD optimizer to take structured inputs

    >>> new_sgd = structured_sgd(sgd)

    Data

    >>> y = np.linspace(1, 10, 100) + np.random.randn(100) + 1
    >>> X = np.array([np.ones(100), np.linspace(1, 100, 100)]).T
    >>> data = np.hstack((y[:, np.newaxis], X))

    Constant Initial values

    >>> w_0 = Parameter(np.array([1., 1.]), Bound())
    >>> lambda_0 = Parameter(.25, Bound())

    >>> res = new_sgd(cost, [w_0, lambda_0], data, batch_size=10,
    ...               eval_obj=True)
    >>> res_w, res_lambda = res.x

    Random Initial values

    >>> from scipy.stats import norm, gamma
    >>> w_0 = Parameter(norm(), Bound(), shape=(2,))
    >>> lambda_0 = Parameter(gamma(1.), Bound())

    >>> res = new_sgd(cost, [w_0, lambda_0], data, batch_size=10,
    ...               eval_obj=True, nstarts=100)
    >>> res_w, res_lambda = res.x
    """
    @wraps(sgd)
    def new_sgd(fun,
                parameters,
                data,
                eval_obj=False,
                batch_size=10,
                args=(),
                random_state=None,
                nstarts=100,
                **sgd_kwargs):

        (array1d, fbounds), shapes = flatten(parameters,
                                             hstack=bt.hstack,
                                             shape=bt.shape,
                                             ravel=bt.ravel
                                             )

        flatten_args_dec = flatten_args(shapes)
        new_fun = flatten_args_dec(fun)

        # Find best random starting candidate if we are doing random starts
        if eval_obj and nstarts > 0:
            data_gen = gen_batch(data, batch_size, random_state=random_state)
            array1d = _random_starts(
                fun=fun,
                parameters=parameters,
                jac=True,
                args=args,
                data_gen=data_gen,
                nstarts=nstarts,
                random_state=random_state
            )

        if bool(eval_obj):
            new_fun = flatten_func_grad(new_fun)
        else:
            new_fun = flatten_grad(new_fun)

        result = sgd(new_fun, array1d, data=data, bounds=fbounds, args=args,
                     eval_obj=eval_obj, random_state=random_state,
                     **sgd_kwargs)

        result['x'] = tuple(unflatten(result['x'], shapes))
        return result

    return new_sgd