import numpy as np
import pandas as pd
from strategies.base import Strategy


class MomentumStrategy(Strategy):
    """
    Classic cross-sectional momentum strategy (Jegadeesh & Titman, 1993).

    Ranks assets by their past (lookback - skip) day return, skipping the most
    recent `skip` days to avoid short-term mean reversion. Allocates equally
    among the top_n winners each period.

    Default parameters approximate the 12-1 month academic factor.
    """

    def __init__(self, lookback: int = 252, skip: int = 21, top_n: int = 5):
        self.lookback = lookback  # ~12 months
        self.skip = skip          # ~1 month — skip to avoid reversal
        self.top_n = top_n

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        # Momentum: return from T-(lookback) to T-skip
        momentum = data.shift(self.skip).pct_change(self.lookback - self.skip)

        n_valid = momentum.notna().sum(axis=1)
        enough_data = n_valid >= self.top_n

        score_rank = momentum.rank(axis=1, ascending=False, na_option="bottom")
        top_mask = (score_rank <= self.top_n) & enough_data.values.reshape(-1, 1)

        raw = top_mask.astype(float)
        row_sums = raw.sum(axis=1).replace(0, np.nan)
        signals = raw.div(row_sums, axis=0).fillna(0.0)
        signals.iloc[:self.lookback] = 0.0

        return signals.shift(1).fillna(0.0)