"""Market-hedged long-only stock portfolio.

Architecture
------------
Long: top-N stock portfolio (factor-selected, inv-vol weighted).
Short hedge: SPY (or SH ETF) at hedge_ratio × long_gross.

This isolates cross-sectional stock-selection alpha while removing the
broad-market beta. Unlike L/S of individual stocks, the short leg is on
the index — so:
  * no meme-stock blowups on short side (GME/AMC type events)
  * borrow cost on SPY is ≈ 5-15 bps annualized (very cheap)
  * no short-dividend pass-through issue (pay SPY div, but that's offset
    by long-side dividends roughly)

Because the long leg is monthly-rebalanced and the short hedge is fixed
at -1.0 × long_gross, total turnover is dominated by the long leg —
similar to V6 long-only.

Output is PIT-safe via terminal `.shift(1)`.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

from strategies.base import Strategy
from strategies.factor_combo import SIGNAL_REGISTRY


class LongHedgedStock(Strategy):
    """Long-only stock momentum hedged with SPY short."""

    def __init__(
        self,
        signal_name: str = "rec_mfilt+deep_upvol",
        top_n: int = 15,
        rebal_freq: int = 21,
        hedge_symbol: str = "SPY",
        hedge_ratio: float = 1.0,
        long_gross: float = 1.0,
        invvol_window: int = 60,
        invvol_floor: float = 0.10,
        invvol_cap: float = 0.20,
        stock_universe: list[str] | None = None,
        # Regime gate: zero out positions when regime_signal < its MA(ma_window)
        regime_gate: bool = False,
        regime_signal: str = "SPY",
        ma_window: int = 200,
    ) -> None:
        if signal_name not in SIGNAL_REGISTRY:
            raise ValueError(f"Unknown signal: {signal_name}")
        self.signal_name = signal_name
        self.signal_func = SIGNAL_REGISTRY[signal_name]
        self.top_n = top_n
        self.rebal_freq = rebal_freq
        self.hedge_symbol = hedge_symbol
        self.hedge_ratio = hedge_ratio
        self.long_gross = long_gross
        self.invvol_window = invvol_window
        self.invvol_floor = invvol_floor
        self.invvol_cap = invvol_cap
        self.stock_universe = stock_universe
        self.regime_gate = regime_gate
        self.regime_signal = regime_signal
        self.ma_window = ma_window

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        if self.hedge_symbol not in data.columns:
            raise ValueError(f"hedge_symbol {self.hedge_symbol!r} missing from panel")
        universe = self.stock_universe or [
            c for c in data.columns if c != self.hedge_symbol
        ]
        universe = [c for c in universe if c in data.columns]

        stock_panel = data[universe]
        sig = self.signal_func(stock_panel)
        rank = sig.rank(axis=1, ascending=False, na_option="bottom")
        n_valid = sig.notna().sum(axis=1)
        enough = n_valid >= self.top_n
        top_mask = (rank <= self.top_n) & enough.values.reshape(-1, 1)

        # Inv-vol weighting within selection
        rets = stock_panel.pct_change(fill_method=None)
        vol = rets.rolling(self.invvol_window,
                           min_periods=self.invvol_window // 2).std() * np.sqrt(252)
        vol_clipped = vol.clip(lower=self.invvol_floor, upper=self.invvol_cap)
        invvol = (1.0 / vol_clipped).where(top_mask, 0.0)
        row_sums = invvol.sum(axis=1).replace(0, np.nan)
        long_w = invvol.div(row_sums, axis=0).fillna(0.0) * self.long_gross

        # Monthly rebalance
        warmup = 252
        rebal_mask = pd.Series(False, index=data.index)
        rebal_idx = list(range(warmup, len(data), self.rebal_freq))
        rebal_mask.iloc[rebal_idx] = True
        long_w[~rebal_mask] = np.nan
        long_w = long_w.ffill().fillna(0.0)
        long_w.iloc[:warmup] = 0.0

        # Build full weights frame: longs in stocks, short in SPY
        out = pd.DataFrame(0.0, index=data.index, columns=data.columns)
        for c in long_w.columns:
            if c in out.columns:
                out[c] = long_w[c]

        # Short hedge: only when long leg is active (gross > 0)
        long_gross_now = long_w.abs().sum(axis=1)
        active = (long_gross_now > 0)
        out[self.hedge_symbol] = -self.hedge_ratio * long_gross_now * active.astype(float)

        # Regime gate: zero everything when regime signal is in bear regime.
        # Avoids the negative-carry case where long stocks tank with SPY but
        # the hedge can't fully offset (since long has higher beta).
        if self.regime_gate and self.regime_signal in data.columns:
            regime_px = data[self.regime_signal]
            ma = regime_px.rolling(self.ma_window).mean()
            risk_on = (regime_px > ma).astype(float).fillna(0.0)
            out = out.mul(risk_on, axis=0)

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


__all__ = ["LongHedgedStock"]