"""Short-horizon momentum alpha."""
import pandas as pd

from pipeline.alpha.base import BaseAlpha
from pipeline.alpha.registry import register_alpha


@register_alpha
class MomentumAlpha(BaseAlpha):
    """Positive trailing return: stocks that rose score high."""

    name = "momentum"

    def __init__(self, lookback: int = 5):
        self.lookback = lookback

    def signal(self, close: pd.DataFrame) -> pd.DataFrame:
        return close.pct_change(self.lookback)