"""
Improved Momentum Quality Strategy.

Improvements over base MomentumQualityStrategy:
1. Monthly rebalancing (original rebalances daily → high turnover)
2. Added recovery factor (strong predictor per IC analysis)
3. Replaced expensive .apply() consistency calc with vectorized version
4. Inverse-vol weighting instead of equal-weight
5. NaN handling fixed throughout
"""

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


class ImprovedMomentumQualityStrategy(Strategy):
    """
    Momentum + quality + recovery with monthly rebal and inv-vol weighting.
    """

    def __init__(
        self,
        momentum_period: int = 252,
        skip: int = 21,
        quality_window: int = 252,
        recovery_window: int = 63,
        vol_window: int = 60,
        rebal_freq: int = 21,
        top_n: int = 20,
    ):
        self.momentum_period = momentum_period
        self.skip = skip
        self.quality_window = quality_window
        self.recovery_window = recovery_window
        self.vol_window = vol_window
        self.rebal_freq = rebal_freq
        self.top_n = top_n

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        # --- Momentum factor ---
        momentum = data.shift(self.skip).pct_change(self.momentum_period - self.skip)

        # --- Quality: return consistency (vectorized) ---
        # Fraction of positive 21-day returns over rolling window
        monthly_ret = data.pct_change(21)
        positive_indicator = (monthly_ret > 0).astype(float)
        consistency = positive_indicator.rolling(
            self.quality_window, min_periods=self.quality_window // 2
        ).mean()

        # --- Quality: inverse max drawdown ---
        rolling_max = data.rolling(self.quality_window, min_periods=self.quality_window // 2).max()
        drawdown = data / rolling_max - 1
        worst_dd = drawdown.rolling(self.quality_window, min_periods=self.quality_window // 2).min()
        inv_dd = -worst_dd  # higher = smaller drawdown = better

        # --- Recovery factor ---
        recovery = data / data.rolling(self.recovery_window, min_periods=self.recovery_window).min() - 1

        # --- Cross-sectional ranking ---
        mom_rank = momentum.rank(axis=1, pct=True, na_option="keep")
        con_rank = consistency.rank(axis=1, pct=True, na_option="keep")
        dd_rank = inv_dd.rank(axis=1, pct=True, na_option="keep")
        rec_rank = recovery.rank(axis=1, pct=True, na_option="keep")

        # Composite: momentum 35%, recovery 25%, consistency 20%, drawdown 20%
        scores = (0.35 * mom_rank + 0.25 * rec_rank +
                  0.20 * con_rank + 0.20 * dd_rank)

        # --- Select top_n ---
        n_valid = scores.notna().sum(axis=1)
        enough = n_valid >= self.top_n
        score_rank = scores.rank(axis=1, ascending=False, na_option="bottom")
        top_mask = (score_rank <= self.top_n) & enough.values.reshape(-1, 1)

        # --- Inverse-vol weighting ---
        returns = data.pct_change()
        vol = returns.rolling(self.vol_window, min_periods=30).std().replace(0, np.nan)
        inv_vol = (1.0 / vol).where(top_mask, 0.0)
        row_sums = inv_vol.sum(axis=1).replace(0, np.nan)
        signals = inv_vol.div(row_sums, axis=0).fillna(0.0)

        # --- Monthly rebalance ---
        warmup = max(self.momentum_period, self.quality_window, self.recovery_window) + self.skip
        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)