# Credits original docs:
# https://www.quantconnect.com/docs/v2/writing-algorithms/scheduled-events
# https://www.quantconnect.com/docs/v2/writing-algorithms/trading-and-orders/position-sizing
# region imports
from AlgorithmImports import *
# endregion


class StrongQqqStockBoostVaaRotation(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100000)
        self.settings.seed_initial_prices = True
        self.settings.free_portfolio_value_percentage = 0.03
        self._total_weight = 0.99
        self._stock_weight = 0.415
        self._weights = [12.0, 4.0, 2.0, 1.0]
        periods = [21, 63, 126, 252]
        self._stocks = "NVDA AVGO AMD MSFT AMZN META GOOGL TSLA NFLX ORCL PANW CRWD SMCI MSTR ANET MU NOW".split() + "UBER AAPL LLY COST ADBE".split()
        self._etfs = "QQQ XLK XLE DBC GLD".split()
        self._defs = "GLD BIL".split()
        self._security_by_ticker = {}
        for ticker in "SPY QQQ XLK XLE DBC GLD BIL".split() + self._stocks:
            security = self.add_equity(ticker, Resolution.DAILY)
            security.momentum = [self.rocp(security, period) for period in periods]
            self._security_by_ticker[ticker] = security
        self.set_warm_up(max(periods) + 20, Resolution.DAILY)
        self.schedule.on(self.date_rules.week_start(self._security_by_ticker["SPY"]), self.time_rules.at(8, 0), self._rebalance)

    def on_warmup_finished(self):
        self._rebalance()

    def _rebalance(self):
        if self.is_warming_up:
            return
        targets = [PortfolioTarget(self._security_by_ticker[ticker], weight) for ticker, weight in self._allocations().items()]
        self.set_holdings(targets, liquidate_existing_holdings=True)

    def _allocations(self):
        # Boost into the strongest momentum stocks when QQQ momentum clears the threshold.
        if self._ready("QQQ") and self._score("QQQ") > 0.35:
            picks = self._rank(self._stocks, 4, True)
            if picks:
                hedge = "GLD" if self._ready("GLD") and self._score("GLD") > 0 else "BIL"
                return {**self._equal(picks, self._stock_weight), hedge: self._total_weight - self._stock_weight}
        # Otherwise rotate into the best trending ETFs while SPY momentum is positive.
        if self._ready("SPY") and self._score("SPY") > 0:
            picks = self._rank(self._etfs, 2, True)
            if picks:
                return self._equal(picks, self._total_weight)
        # Fall back to the best defensive asset, defaulting to T-bills.
        picks = self._rank(self._defs, 1, True) or self._rank(self._defs, 1, False) or ["BIL"]
        return self._equal(picks, self._total_weight)

    def _rank(self, tickers, count, require_positive):
        score_by_ticker = {}
        for ticker in tickers:
            if self._ready(ticker):
                score = self._score(ticker)
                if score > 0 or not require_positive:
                    score_by_ticker[ticker] = score
        return sorted(score_by_ticker, key=lambda ticker: score_by_ticker[ticker])[-count:]

    def _equal(self, tickers, total_weight):
        return {ticker: total_weight / len(tickers) for ticker in tickers}

    def _ready(self, ticker):
        return all(indicator.is_ready for indicator in self._security_by_ticker[ticker].momentum)

    def _score(self, ticker):
        return sum(weight * indicator.current.value for weight, indicator in zip(self._weights, self._security_by_ticker[ticker].momentum))