import numpy as np
import pandas as pd
from strategies.base import StockStrategy, mask_scores


class InverseVolatilityStrategy(StockStrategy):
    """
    Risk parity via inverse-volatility weighting.

    Allocates capital inversely proportional to each asset's realized volatility
    over a rolling window. Assets with lower recent volatility receive larger
    allocations, equalizing the risk contribution of each position.

    This is a simple but robust baseline for risk-based portfolio construction.
    """

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

    def generate_signals(self, data: pd.DataFrame, eligible: pd.DataFrame | None = None) -> pd.DataFrame:
        returns = data.pct_change(fill_method=None)
        vol = returns.rolling(self.vol_window).std()

        # Replace zero vol with NaN to avoid division by zero
        vol = vol.replace(0, np.nan)
        inv_vol = mask_scores((1.0 / vol).where(data.notna()), eligible)

        row_sums = inv_vol.sum(axis=1).replace(0, np.nan)
        signals = inv_vol.div(row_sums, axis=0).fillna(0.0)
        signals.iloc[:self.vol_window] = 0.0

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