"""
Cross-asset time-series momentum strategy (ETF-only, inherently PIT).

Alpha source: Moskowitz, Ooi, Pedersen (2012) — assets with positive
12-month returns continue to trend. Earns during equity crises when
bonds/gold trend up while stocks trend down.

Universe: SPY, GLD, TLT, IEF, DBC (broad, liquid ETFs)
Signal: 12-month total return; go long top K assets with positive momentum
Rebalance: monthly (21 trading days)
If no asset has positive 12m return → 100% cash (SHY proxy = 0 weights)
"""

import numpy as np
import pandas as pd


class CrossAssetMomentum:
    """Time-series momentum across major asset classes."""

    UNIVERSE = ["SPY", "GLD", "TLT", "IEF", "DBC"]

    def __init__(
        self,
        lookback: int = 252,
        top_k: int = 3,
        rebal_freq: int = 21,
        vol_scale: bool = True,
        vol_window: int = 63,
    ):
        self.lookback = lookback
        self.top_k = top_k
        self.rebal_freq = rebal_freq
        self.vol_scale = vol_scale
        self.vol_window = vol_window

    def generate_signals(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        Parameters
        ----------
        data : DataFrame with columns including UNIVERSE ETFs (prices).

        Returns
        -------
        DataFrame of weights aligned to data.index and data.columns.
        """
        # Restrict to our universe (ignore missing gracefully)
        available = [t for t in self.UNIVERSE if t in data.columns]
        prices = data[available]

        # 12-month return signal (shifted 1 day for execution lag)
        mom = prices.pct_change(self.lookback).shift(1)

        # Inverse-vol for position sizing (optional)
        if self.vol_scale:
            daily_ret = prices.pct_change()
            vol = daily_ret.rolling(self.vol_window).std() * np.sqrt(252)
            inv_vol = (1.0 / vol).replace([np.inf, -np.inf], np.nan)
        else:
            inv_vol = None

        weights = pd.DataFrame(0.0, index=data.index, columns=data.columns)
        n = len(prices.index)
        last_w = pd.Series(0.0, index=available)

        for i in range(self.lookback + 1, n, self.rebal_freq):
            row_mom = mom.iloc[i]
            # Only go long assets with positive momentum
            positive = row_mom[row_mom > 0].sort_values(ascending=False)

            if positive.empty:
                last_w = pd.Series(0.0, index=available)
            else:
                selected = positive.head(self.top_k).index.tolist()

                if self.vol_scale and inv_vol is not None:
                    iv = inv_vol.iloc[i][selected]
                    if iv.sum() > 0:
                        w = iv / iv.sum()
                    else:
                        w = pd.Series(1.0 / len(selected), index=selected)
                else:
                    w = pd.Series(1.0 / len(selected), index=selected)

                last_w = pd.Series(0.0, index=available)
                last_w[selected] = w.values

            # Hold until next rebalance
            end_i = min(i + self.rebal_freq, n)
            for col in available:
                weights.iloc[i:end_i, weights.columns.get_loc(col)] = last_w[col]

        return weights