# region imports
from AlgorithmImports import *
# endregion

class DualMomentumRotation(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)
        self.SetEndDate(2022, 12, 31)
        self.SetCash(100000)
        self.SetBrokerageModel(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        # Universe: US (SPY), Developed ex-US (EFA), Emerging (EEM) with bond fallback (IEF)
        self.risk_assets = ["SPY", "EFA", "EEM"]
        self.safe_asset = "IEF"
        self.symbols = {t: self.AddEquity(t, Resolution.Daily).Symbol for t in self.risk_assets + [self.safe_asset]}

        # 12-month momentum, evaluated monthly at month end
        self.lookback_months = 12
        self.SetWarmUp(self.lookback_months * 31, Resolution.Daily)

        # Schedule monthly on last trading day, after market open for stability
        self.Schedule.On(self.DateRules.MonthEnd(self.symbols[self.risk_assets[0]], 0),
                         self.TimeRules.AfterMarketOpen(self.symbols[self.risk_assets[0]], 30),
                         self.Rebalance)

        # Absolute momentum threshold vs. cash/bonds
        self.absolute_momentum_threshold = 0.0

    def Rebalance(self):
        if self.IsWarmingUp:
            return

        # Pull multi-symbol history once to avoid per-ticker calls
        bars = self.History(list(self.symbols.values()), self.lookback_months * 31, Resolution.Daily)
        if bars is None or bars.empty:
            return

        rets = {}
        for ticker in self.risk_assets + [self.safe_asset]:
            sym = self.symbols[ticker]
            df = None
            try:
                df = bars.xs(sym, level=0)
            except Exception:
                # If not multi-indexed, try filtering by column
                if isinstance(bars, pd.DataFrame) and 'symbol' in bars.columns:
                    df = bars[bars['symbol'] == sym]
                else:
                    df = bars
            if df is None or df.empty:
                continue
            series = df['close'] if 'close' in df else df.Close
            monthly = series.resample('M').last()
            if len(monthly) < self.lookback_months + 1:
                continue
            past = monthly.iloc[-(self.lookback_months+1)]
            now = monthly.iloc[-1]
            if past and past != 0:
                rets[ticker] = (now / past) - 1.0

        # Need all returns computed
        if set(rets.keys()) != set(self.risk_assets + [self.safe_asset]):
            return

        ranked = sorted(self.risk_assets, key=lambda t: rets[t], reverse=True)
        top = ranked[0]
        target = top if rets[top] > self.absolute_momentum_threshold else self.safe_asset

        targets = []
        for t in self.risk_assets + [self.safe_asset]:
            sym = self.symbols[t]
            weight = 1.0 if t == target else 0.0
            targets.append(PortfolioTarget(sym, weight))

        self.Debug(f"{self.Time.date()} Momentum: " + ", ".join([f"{k}:{rets[k]:.2%}" for k in rets]))
        self.SetHoldings(targets)

    def OnSecuritiesChanged(self, changes):
        pass