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


class MeanReversionStrategy(Strategy):
    """
    Monthly "buy the dip" with momentum confirmation.

    Among stocks with positive 12-month momentum, overweight those
    that dipped most in the past month. Rebalances monthly to keep
    turnover low. Combines long-term trend (avoid losers) with
    short-term mean reversion (buy winners on sale).
    """

    def __init__(self, mom_lookback: int = 252, dip_window: int = 21,
                 rebal_freq: int = 21, top_n: int = 20):
        self.mom_lookback = mom_lookback
        self.dip_window = dip_window
        self.rebal_freq = rebal_freq
        self.top_n = top_n

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        long_mom = data.pct_change(self.mom_lookback)
        has_momentum = long_mom > 0

        short_ret = data.pct_change(self.dip_window)

        # Among positive-momentum stocks, rank by biggest dip
        scores = short_ret.where(has_momentum, np.nan)

        # Rank: most negative = rank 1 (biggest dip)
        rank = scores.rank(axis=1, ascending=True, na_option="bottom")
        n_valid = scores.notna().sum(axis=1)
        effective_n = n_valid.clip(upper=self.top_n)

        top_mask = rank <= effective_n.values.reshape(-1, 1)
        top_mask = top_mask & has_momentum  # ensure we only pick momentum stocks

        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)

        warmup = self.mom_lookback + self.dip_window

        # Only keep rebalance-day signals, forward-fill between them
        rebal_mask = pd.Series(False, index=data.index)
        rebal_indices = range(warmup, len(data), self.rebal_freq)
        rebal_mask.iloc[list(rebal_indices)] = True

        # Zero out non-rebalance days then forward-fill
        signals[~rebal_mask] = np.nan
        signals = signals.ffill().fillna(0.0)

        signals.iloc[:warmup] = 0.0
        return signals.shift(1).fillna(0.0)