# region imports
from AlgorithmImports import *
# endregion


class Top2EtfRotator(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100_000)
        self.settings.free_portfolio_value_percentage = 0.05
        self._lookback_6m = 126
        self._lookback_12m = 252
        self._securities = []
        self._safe_bond_security = None
        self._spy_security = None
        for ticker in ["SPY", "QQQ", "EFA", "IEF", "TLT", "GLD", "VNQ", "DBC", "HYG", "XLU"]:
            security = self.add_equity(ticker, Resolution.DAILY)
            security.set_data_normalization_mode(DataNormalizationMode.TOTAL_RETURN)
            self._securities.append(security)
            if ticker == "IEF":
                self._safe_bond_security = security
            if ticker == "SPY":
                self._spy_security = security
        if self._safe_bond_security is None:
            self._safe_bond_security = self._securities[0]
        if self._spy_security is None:
            self._spy_security = self._securities[0]
        # Only these three risk-on ETFs are eligible for the momentum rotation.
        self._risk_on_securities = [s for s in self._securities if s.symbol.value in ["SPY", "QQQ", "GLD"]]
        if not self._risk_on_securities:
            self._risk_on_securities = [s for s in self._securities if s != self._safe_bond_security]
        self.set_warm_up(self._lookback_12m + 10)
        self.schedule.on(self.date_rules.week_start(self._securities[0]), self.time_rules.at(8, 0), self._rebalance)

    def on_warmup_finished(self):
        self._rebalance()

    def _rebalance(self):
        if self.is_warming_up or not self._safe_bond_security.has_data or self._safe_bond_security.is_delisted:
            return
        max_lookback = self._lookback_12m + 1
        history = self.history([s.symbol for s in self._securities], max_lookback, Resolution.DAILY)
        if history.empty:
            self._go_full_safe_bond()
            return
        index_level = history.index.get_level_values(0).unique()
        momentum_scores = {}
        # Use SPY's own 6-month momentum as an absolute-momentum regime gate.
        spy_closes = history.loc[self._spy_security.symbol]["close"].astype(float) if self._spy_security.symbol in index_level else None
        if spy_closes is None or len(spy_closes) < self._lookback_6m + 1 or (spy_closes.iloc[-1] / spy_closes.iloc[-self._lookback_6m - 1]) - 1.0 <= 0.0:
            self._go_full_safe_bond()
            return
        for security in self._risk_on_securities:
            if security.symbol not in index_level:
                continue
            closes = history.loc[security.symbol]["close"].astype(float)
            if len(closes) < max_lookback or len(closes) < self._lookback_12m + 1 or len(closes) < self._lookback_6m + 1:
                continue
            # Score each ETF by an equal blend of 6-month and 12-month momentum.
            momentum_scores[security] = 0.5 * ((closes.iloc[-1] / closes.iloc[-self._lookback_6m - 1]) - 1.0) + 0.5 * ((closes.iloc[-1] / closes.iloc[-self._lookback_12m - 1]) - 1.0)
        if not momentum_scores:
            self._go_full_safe_bond()
            return
        ranked = sorted(momentum_scores, key=lambda k: momentum_scores[k], reverse=True)
        if momentum_scores[ranked[0]] <= 0.0:
            self._go_full_safe_bond()
            return
        # Concentrate the whole book in the single strongest ETF, leaving the safe bond flat.
        top_securities = ranked[:1]
        target_weights = {s: 1.0 / len(top_securities) for s in top_securities}
        target_weights[self._safe_bond_security] = 0.0
        total_weight = sum(target_weights.values())
        if total_weight > 0:
            for security in target_weights:
                target_weights[security] /= total_weight
        targets = [PortfolioTarget(security.symbol, target_weights.get(security, 0.0)) for security in self._securities]
        self.set_holdings(targets, liquidate_existing_holdings=True)

    def _go_full_safe_bond(self):
        targets = [PortfolioTarget(security.symbol, 1.0 if security == self._safe_bond_security else 0.0) for security in self._securities]
        self.set_holdings(targets, liquidate_existing_holdings=True)