def next_batch_for_narx(dataset, window_size, n_steps, batch_size):
    """Returns a batch in a narx format: x0, y, inputs.
    Args:
        dataset (TSDataset): time series dataset.
        window_size (int): number of steps to include in x0.
        n_steps (int): number of steps to include in y and u.
        batch_size (int): number of samples in the batch.
    Returns:
        (x0, y, u): batch.
    """
    def format_batch(batch, window_size, n_steps):
        _y = [batch['y'][i, :, :]
              for i in range(batch['y'].shape[0])]
        x0 = _y[:window_size]
        y = _y[window_size:]
        # dy = [_y[i] - _y[i - 1] for i in range(window_size, len(_y))]
        u = [batch['inputs'][i, :, :]
             for i in range(window_size - 1, batch['inputs'].shape[0] - 1)]

        assert (len(x0) == window_size and
                len(y) == n_steps and
                len(u) == n_steps), \
            'len of x0, y or u do not coincide with provided '\
            'window_size and n_steps'
        return x0, y, u
    batch = dataset.next_batch(window_size=window_size + n_steps,
                               batch_size=batch_size)
    x0, y, u = format_batch(batch, window_size, n_steps)
    return x0, y, u