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


class MultiFactorStrategy(Strategy):
    """
    Multi-factor strategy combining momentum and value signals with a market-timing filter.

    Factors:
    - Momentum: past return from (momentum_period + skip) to skip days ago (avoids short-term reversal)
    - Value: rolling min / current price (inverted price-to-low ratio — cheaper = higher score)
    - Market timing: only invest when SPY is above its long-term moving average

    Signal generation is fully vectorized — no Python loops over time.
    """

    def __init__(
        self,
        tickers,
        benchmark: str = "SPY",
        window: int = 200,
        momentum_period: int = 230,
        momentum_skip: int = 20,
        value_period: int = 250,
        top_n: int = 5,
    ):
        self.tickers = list(tickers)
        self.benchmark = benchmark
        self.window = window
        self.momentum_period = momentum_period
        self.momentum_skip = momentum_skip
        self.value_period = value_period
        self.top_n = top_n

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        stock = data[self.tickers]

        # --- Market timing filter ---
        spy_ma = data[self.benchmark].rolling(self.window).mean()
        market_up = (data[self.benchmark] > spy_ma).values  # shape (T,)

        # --- Momentum factor ---
        # Return from T-(momentum_period+skip) to T-skip, avoiding the last month
        momentum = stock.shift(self.momentum_skip).pct_change(self.momentum_period)

        # --- Value factor ---
        # min_price_over_period / current_price  (higher = more "undervalued" vs recent range)
        value = stock.rolling(self.value_period).min() / stock

        # --- Cross-sectional ranking (each row ranked across assets) ---
        mom_rank = momentum.rank(axis=1, pct=True, na_option="bottom")
        val_rank = value.rank(axis=1, pct=True, na_option="bottom")
        scores = mom_rank + val_rank  # combined score, higher = better

        # --- Select top_n assets per row ---
        # Only allocate rows that have enough valid scores
        n_valid = scores.notna().sum(axis=1)
        enough_data = n_valid >= self.top_n

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

        # Equal-weight allocation among selected assets
        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)

        # --- Apply market timing: zero out when SPY is below its MA ---
        signals[~market_up] = 0.0

        # --- Zero out warm-up period ---
        warmup = max(self.window, self.momentum_period + self.momentum_skip, self.value_period)
        signals.iloc[:warmup] = 0.0

        # Shift by 1: signal computed at close of day t trades at open of day t+1
        return signals.shift(1).fillna(0.0)