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


class DualMomentumStrategy(Strategy):
    """
    Dual momentum with multi-timeframe confirmation.

    Unlike plain momentum (single 12-1m lookback), this requires agreement
    across THREE timeframes — short (1-3m), medium (3-6m), and long (6-12m).
    A stock must show positive returns on ALL three to qualify.

    This filters out stocks riding a single spike and favors sustained trends.
    Position count is variable: when few stocks pass all three filters,
    the strategy naturally goes to partial cash.
    """

    def __init__(self, top_n: int = 20):
        self.top_n = top_n

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        # Three timeframe returns (each skipping most recent 5 days)
        skip = 5
        short_mom = data.shift(skip).pct_change(63 - skip)    # ~3 month
        med_mom = data.shift(skip).pct_change(126 - skip)     # ~6 month
        long_mom = data.shift(skip).pct_change(252 - skip)    # ~12 month

        # All three must be positive (absolute momentum across timeframes)
        all_positive = (short_mom > 0) & (med_mom > 0) & (long_mom > 0)

        # Composite score: average percentile rank across timeframes
        short_rank = short_mom.rank(axis=1, pct=True, na_option="bottom")
        med_rank = med_mom.rank(axis=1, pct=True, na_option="bottom")
        long_rank = long_mom.rank(axis=1, pct=True, na_option="bottom")
        composite = (short_rank + med_rank + long_rank) / 3

        # Only consider stocks passing absolute filter
        composite_filtered = composite.where(all_positive, np.nan)

        n_valid = composite_filtered.notna().sum(axis=1)
        enough = n_valid >= 1

        rank = composite_filtered.rank(axis=1, ascending=False, na_option="bottom")
        effective_n = n_valid.clip(upper=self.top_n)
        top_mask = (rank <= effective_n.values.reshape(-1, 1)) & enough.values.reshape(-1, 1)
        top_mask = top_mask & all_positive

        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[:252] = 0.0

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