# https://quantpedia.com/strategies/refining-etf-asset-momentum-strategy/
# 
# The investment universe comprises 13 Exchange-Traded Funds (ETFs) from diverse asset classes, including 6 stock ETFs, 3 bond ETFs, 3 commodity ETFs, 
# and 1 currency ETF, traded between April 10, 2006, and February 28, 2023. Adjusted close price data is sourced from Yahoo Finance to account for dividends, 
# splits, and other corporate actions. Momentum is calculated over 3, 6, 9, and 12-month periods, ranking ETFs by performance. The top 4 ETFs are selected 
# for equally weighted long positions, while the single worst-performing ETF is chosen for a 30% weighted short position. A correlation filter determines 
# market conditions: when short-term (20-day) correlation exceeds long-term (250-day) correlation, only long positions are maintained; when short-term 
# correlation is lower, the long+short positions are applied. The portfolio is rebalanced monthly.

# region imports
from AlgorithmImports import *
# endregion

class RefiningETFAssetMomentumStrategy(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2006, 4, 1)
        self.set_cash(100_000)
        
        self._long_period: int = 250
        self._short_period: int = 20
        self._month_period: int = 21
        self._long_count: int = 4
        self._short_count: int = 1
        self._short_weight: float = .3
        self._mom_periods: List[int] = [3, 6, 9, 12]

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        tickers: List[str] = ["SPY", "IWM", "EFA", "EEM", "IYR", "QQQ", "LQD", "IEF", "TIP", "GLD", "USO", "DBC", "FXE"]
        self._traded_assets: List[Symbol] = [
            self.add_equity(ticker, Resolution.DAILY).symbol for ticker in tickers
        ]

        self._selection_flag: bool = False
        self.settings.minimum_order_margin_portfolio_percentage = 0.
        self.settings.daily_precise_end_time = False

        self.schedule.on(self.date_rules.month_end(self._traded_assets[0]), 
                        self.time_rules.before_market_close(self._traded_assets[0]), 
                        self.selection)

    def on_data(self, slice: Slice) -> None:
        # monthly rebalance
        if not self._selection_flag:
            return
        self._selection_flag = False

        long: List[str] = []
        short: List[str] = []

        # correlation signal
        history: DataFrame = self.history(
            TradeBar, self._traded_assets, max(self._mom_periods) * self._month_period
        )

        prices: DataFrame = history.close.unstack(level=0).dropna()

        if len(prices) != max(self._mom_periods) * self._month_period or len(prices.columns) != len(self._traded_assets):
            self.log('Not enough data for signal calculation.')
            return

        returns: DataFrame = prices.pct_change().dropna()
        long_corr: DataFrame = returns.iloc[-self._long_period:].corr()
        short_corr: DataFrame = returns.iloc[-self._short_period:].corr()

        long_corr_mean: float = long_corr.values[np.triu_indices_from(long_corr.values, 1)].mean()
        short_corr_mean: float = short_corr.values[np.triu_indices_from(short_corr.values, 1)].mean()

        trade_directions: bool = (1, 0) if short_corr_mean > long_corr_mean else (1, -self._short_weight)

        mom_avg: Dict[Symbol] = {}

        # average momentum calculation and sort
        for column in prices.columns:
            mom_avg[column] = np.mean([
                prices[column].iloc[-1] / prices[column].iloc[-period * self._month_period] - 1 
                for period in self._mom_periods
            ])

        sorted_mom_avg: List[Symbol] = sorted(mom_avg, key=mom_avg.get, reverse=True)
        long: List[Symbol] = sorted_mom_avg[:self._long_count]
        short: List[Symbol] = sorted_mom_avg[-self._short_count:]
 
        # order execution
        targets: List[PortfolioTarget] = []
        for i, portfolio in enumerate([long, short]):
            for symbol in portfolio:
                if slice.contains_key(symbol) and slice[symbol]:
                    targets.append(PortfolioTarget(symbol, trade_directions[i] / len(portfolio)))
    
        self.set_holdings(targets, True)

    def selection(self) -> None:
        self._selection_flag = True