import numpy as np

def moving_average_rewards(rewards, window_size=10):
    """
    Computes a moving average of episodic rewards for smoother visualization.

    Parameters
    ----------
    rewards : list or np.ndarray
        Sequence of total rewards per episode.
    window_size : int, optional (default=10)
        Number of episodes to average over.

    Returns
    -------
    np.ndarray
        Smoothed rewards.

    Example
    -------
    >>> rewards = [10, 20, 30, 40, 50]
    >>> moving_average_rewards(rewards, window_size=3)
    array([10., 15., 20., 30., 40.])
    """
    if not isinstance(rewards, (list, np.ndarray)):
        raise TypeError("rewards must be a list or numpy array")

    rewards = np.array(rewards, dtype=float)
    if window_size <= 0:
        raise ValueError("window_size must be positive")

    cumsum = np.cumsum(np.insert(rewards, 0, 0))
    moving_avg = (cumsum[window_size:] - cumsum[:-window_size]) / window_size

    # Pad the beginning so the output matches input length
    padding = np.ones(window_size - 1) * moving_avg[0] if len(moving_avg) > 0 else []
    return np.concatenate([padding, moving_avg])


if __name__ == "__main__":
    rewards = [10, 20, 30, 40, 50]
    print("Moving average:", moving_average_rewards(rewards, window_size=3))
