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


class RecoveryMomentumStrategy(Strategy):
    """
    Recovery + momentum composite factor strategy.

    Combines the two strongest predictive factors (by IC analysis):
    1. Recovery: price relative to 63-day low (how far stock has bounced
       from its recent trough). Higher = stronger recovery trajectory.
    2. Long-term momentum: 12-1 month return. Avoids most recent month
       to sidestep short-term reversal noise.

    Rank-combines both factors (50/50), selects top_n, equal-weights.
    Rebalances monthly to control turnover.

    The recovery factor captures stocks in strong V-shaped rebounds —
    a pattern that tends to persist. Combined with momentum, it filters
    for stocks with both long-term strength AND recent acceleration.
    """

    def __init__(self, recovery_window: int = 63, mom_lookback: int = 252,
                 mom_skip: int = 21, rebal_freq: int = 21, top_n: int = 10):
        self.recovery_window = recovery_window
        self.mom_lookback = mom_lookback
        self.mom_skip = mom_skip
        self.rebal_freq = rebal_freq
        self.top_n = top_n

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        # Factor 1: Recovery — price / rolling min
        recovery = data / data.rolling(self.recovery_window).min() - 1

        # Factor 2: 12-1 month momentum
        momentum = data.shift(self.mom_skip).pct_change(self.mom_lookback - self.mom_skip)

        # Cross-sectional percentile ranks
        rec_rank = recovery.rank(axis=1, pct=True, na_option="bottom")
        mom_rank = momentum.rank(axis=1, pct=True, na_option="bottom")

        # Composite score (50/50)
        composite = 0.5 * rec_rank + 0.5 * mom_rank

        # Select top_n
        rank = composite.rank(axis=1, ascending=False, na_option="bottom")
        n_valid = composite.notna().sum(axis=1)
        enough = n_valid >= self.top_n
        top_mask = (rank <= self.top_n) & enough.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)

        # Monthly rebalance: keep only rebal-day signals, forward-fill
        warmup = max(self.mom_lookback, self.recovery_window)
        rebal_mask = pd.Series(False, index=data.index)
        rebal_indices = list(range(warmup, len(data), self.rebal_freq))
        rebal_mask.iloc[rebal_indices] = True

        signals[~rebal_mask] = np.nan
        signals = signals.ffill().fillna(0.0)
        signals.iloc[:warmup] = 0.0

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