import numpy as np

def n_step_return(rewards, gamma=0.99, n=None, next_state_value=0.0):
    """
    Calculates the n-step return for a sequence of rewards.

    Parameters
    ----------
    rewards : list or np.ndarray
        A list of rewards [r_t, r_{t+1}, ..., r_{t+n-1}].
    gamma : float, optional (default=0.99)
        Discount factor (0 <= gamma <= 1).
    n : int, optional
        Number of steps to include. If None, uses all rewards.
    next_state_value : float, optional (default=0.0)
        The estimated value of the next state V(s_{t+n}).

    Returns
    -------
    float
        The computed n-step return.

    Example
    -------
    >>> rewards = [1, 0, 2]
    >>> n_step_return(rewards, gamma=0.9, n=3, next_state_value=1.0)
    1 + 0.9*0 + 0.9^2*2 + 0.9^3*1 = 3.439

    Notes
    -----
    This function is commonly used in n-step TD and actor-critic algorithms.
    """
    if not isinstance(rewards, (list, np.ndarray)):
        raise TypeError("rewards must be a list or numpy array")

    if not 0 <= gamma <= 1:
        raise ValueError("gamma must be between 0 and 1")

    rewards = np.array(rewards, dtype=float)
    n = len(rewards) if n is None else min(n, len(rewards))

    discounted_sum = 0.0
    for i in range(n):
        discounted_sum += (gamma ** i) * rewards[i]

    # Add the bootstrapped next-state value
    discounted_sum += (gamma ** n) * next_state_value
    return float(discounted_sum)


if __name__ == "__main__":
    # Demo
    rewards = [1, 0, 2]
    G = n_step_return(rewards, gamma=0.9, n=3, next_state_value=1.0)
    print(f"N-step return: {G:.4f}")
