"""
Composite Alpha Strategy.

Combines the strongest alpha factors discovered in research:
1. Recovery (63d) - strongest single IC
2. Intermediate momentum (7m) - strong trend signal
3. Quality (consistency + low drawdown) - filters lottery tickets
4. MA200 trend confirmation - only stocks above their MA200

With:
- Inverse-vol weighting for risk parity among selected stocks
- SPY MA200 market regime gate (reduce exposure in bear markets)
- Biweekly rebalancing (compromise between signal freshness and turnover)
"""

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


class CompositeAlphaStrategy(Strategy):
    """
    Multi-factor alpha composite with regime gating.
    """

    def __init__(
        self,
        tickers: list[str] | None = None,
        benchmark: str = "SPY",
        recovery_window: int = 63,
        intermediate_period: int = 147,
        skip: int = 21,
        quality_window: int = 252,
        vol_window: int = 60,
        rebal_freq: int = 10,
        top_n: int = 20,
        regime_gate: bool = True,
        regime_ma: int = 200,
    ):
        self.tickers = tickers
        self.benchmark = benchmark
        self.recovery_window = recovery_window
        self.intermediate_period = intermediate_period
        self.skip = skip
        self.quality_window = quality_window
        self.vol_window = vol_window
        self.rebal_freq = rebal_freq
        self.top_n = top_n
        self.regime_gate = regime_gate
        self.regime_ma = regime_ma

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        # Separate benchmark from stocks if needed
        if self.tickers is not None:
            stock_data = data[[t for t in self.tickers if t in data.columns]]
        else:
            stock_cols = [c for c in data.columns if c != self.benchmark]
            stock_data = data[stock_cols]

        # --- Factor 1: Recovery ---
        recovery = stock_data / stock_data.rolling(
            self.recovery_window, min_periods=self.recovery_window
        ).min() - 1

        # --- Factor 2: Intermediate momentum (7-month, skip 1 month) ---
        int_mom = stock_data.shift(self.skip).pct_change(self.intermediate_period - self.skip)

        # --- Factor 3: Quality composite ---
        # Consistency: fraction of positive 21-day returns
        monthly_ret = stock_data.pct_change(21)
        consistency = (monthly_ret > 0).astype(float).rolling(
            self.quality_window, min_periods=self.quality_window // 2
        ).mean()

        # Up-volume proxy: sum of positive daily returns over 20 days
        daily_ret = stock_data.pct_change()
        up_vol_proxy = daily_ret.where(daily_ret > 0, 0).rolling(20, min_periods=15).sum()

        # --- Factor 4: Above MA200 (per-stock trend filter) ---
        ma200 = stock_data.rolling(200, min_periods=200).mean()
        above_ma = (stock_data > ma200)

        # --- Cross-sectional ranks ---
        rec_rank = recovery.rank(axis=1, pct=True, na_option="keep")
        mom_rank = int_mom.rank(axis=1, pct=True, na_option="keep")
        con_rank = consistency.rank(axis=1, pct=True, na_option="keep")
        upv_rank = up_vol_proxy.rank(axis=1, pct=True, na_option="keep")

        # Composite: recovery 30%, int_momentum 30%, consistency 20%, up_volume 20%
        composite = (0.30 * rec_rank + 0.30 * mom_rank +
                     0.20 * con_rank + 0.20 * upv_rank)

        # Apply per-stock MA200 filter: must be in uptrend
        composite = composite.where(above_ma, np.nan)

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

        # --- Inverse-vol weighting ---
        vol = daily_ret.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)

        # --- Market regime gate (SPY > MA200) ---
        if self.regime_gate and self.benchmark in data.columns:
            spy = data[self.benchmark]
            spy_ma = spy.rolling(self.regime_ma, min_periods=self.regime_ma).mean()
            market_bull = (spy > spy_ma).astype(float)
            # Partial scaling: 100% when bullish, 30% when bearish (don't go fully to cash)
            regime_scale = market_bull * 0.7 + 0.3
            signals = signals.mul(regime_scale, axis=0)

        # --- Biweekly rebalance ---
        warmup = max(self.quality_window, 200, self.intermediate_period) + 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

        # Align to full data columns (in case benchmark is in data)
        full_signals = pd.DataFrame(0.0, index=data.index, columns=data.columns)
        for col in signals.columns:
            if col in full_signals.columns:
                full_signals[col] = signals[col]

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