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


class AdaptiveMomentumStrategy(Strategy):
    """
    Momentum with dynamic position sizing via inverse-volatility weighting.

    Combines two proven ideas:
    1. Momentum selection: pick top_n stocks by 12-1 month return
    2. Inverse-vol weighting: instead of equal-weight, allocate more
       to lower-vol winners (smoother ride, better risk-adjusted returns)
    3. Regime scaling: scale total exposure by inverse of market vol
       (reduce exposure in high-vol periods, increase in calm periods)

    This addresses momentum's main weakness: concentration in high-vol
    names that crash hard in risk-off episodes.
    """

    def __init__(self, lookback: int = 252, skip: int = 21,
                 vol_window: int = 60, top_n: int = 20):
        self.lookback = lookback
        self.skip = skip
        self.vol_window = vol_window
        self.top_n = top_n

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

        n_valid = momentum.notna().sum(axis=1)
        enough = n_valid >= self.top_n

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

        # Inverse-vol weighting among selected stocks
        returns = data.pct_change()
        vol = returns.rolling(self.vol_window).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)

        # Regime scaling: compare current market vol to its 1-year median
        market_vol = vol.mean(axis=1)
        vol_median = market_vol.rolling(252).median()
        vol_scale = (vol_median / market_vol).clip(0.3, 1.5)
        signals = signals.mul(vol_scale, axis=0)

        # Re-normalize so max total weight = 1.0
        row_totals = signals.sum(axis=1)
        overflow = row_totals > 1.0
        signals[overflow] = signals[overflow].div(row_totals[overflow], axis=0)

        warmup = max(self.lookback, self.vol_window + 252)
        signals.iloc[:warmup] = 0.0

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