# region imports
from AlgorithmImports import *
# endregion


class VaaGoblinOmega(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100000)
        growth_tickers = ["NVDA", "AVGO", "MSFT", "META", "LLY", "COST", "QQQ", "XLK", "SMH", "SPY"]
        canary_tickers = ["SPY", "QQQ", "IWM", "HYG", "EFA"]
        high_beta_tickers = ["NVDA", "AVGO", "META", "QQQ", "XLK", "SMH"]
        tickers = ["NVDA", "AVGO", "MSFT", "META", "LLY", "COST", "QQQ", "XLK", "SMH", "SPY", "IWM", "HYG", "EFA", "GLD", "SHY", "IEF", "XLV", "XLP", "XLU"]
        self._securities = []
        self._growth_securities = []
        self._canary_securities = []
        self._security_by_ticker = {}
        for ticker in tickers:
            security = self.add_equity(ticker, Resolution.DAILY)
            security.sma_50 = self.sma(security, 50)
            security.sma_100 = self.sma(security, 100)
            security.sma_200 = self.sma(security, 200)
            security.rocp_21 = self.rocp(security, 21)
            security.rocp_63 = self.rocp(security, 63)
            security.rocp_126 = self.rocp(security, 126)
            security.rocp_252 = self.rocp(security, 252)
            security.std_63 = self.std(security, 63)
            security.is_high_beta = ticker in high_beta_tickers
            self._securities.append(security)
            self._security_by_ticker[ticker] = security
            if ticker in growth_tickers:
                self._growth_securities.append(security)
            if ticker in canary_tickers:
                self._canary_securities.append(security)
        self._spy = self._security_by_ticker["SPY"]
        self._qqq = self._security_by_ticker["QQQ"]
        self.set_benchmark(self._spy)
        self.set_warm_up(300, Resolution.DAILY)
        self.schedule.on(self.date_rules.month_start(self._spy, 5), self.time_rules.at(8, 0), self._rebalance)

    def on_warmup_finished(self):
        self._rebalance()

    def _rebalance(self):
        if self.is_warming_up:
            return
        # Skip until every indicator is ready and every security has a valid price.
        for security in self._securities:
            if not (security.sma_50.is_ready and security.sma_100.is_ready and security.sma_200.is_ready and security.rocp_21.is_ready and security.rocp_63.is_ready and security.rocp_126.is_ready and security.rocp_252.is_ready and security.std_63.is_ready):
                return
            if not security.has_data or security.price <= 0:
                return
        # Count canary assets that are both rising and above their long-term trend.
        risk_score = sum(1 for security in self._canary_securities if self._vaa_score(security) > 0 and security.price > security.sma_200.current.value)
        # Pick the allocation profile from the stress check first, then the canary risk score.
        if (self._spy.rocp_21.current.value < -0.045 or self._qqq.rocp_21.current.value < -0.065) and self._spy.price < self._spy.sma_50.current.value and self._qqq.price < self._qqq.sma_50.current.value:
            weights = self._risk_off_weights()
        elif risk_score >= 4:
            weights = self._strong_risk_on_weights()
        elif risk_score >= 3:
            weights = self._weak_risk_on_weights()
        else:
            weights = self._risk_off_weights()
        targets = [PortfolioTarget(security, weights.get(security, 0)) for security in self._securities]
        self.set_holdings(targets, liquidate_existing_holdings=True)

    def _strong_risk_on_weights(self):
        selected = self._select_top_growth_securities(4, 3)
        weights = {}
        for security, weight in zip(selected, [0.36, 0.25, 0.17, 0.10]):
            self._add_weight(weights, security, weight)
        self._add_weight(weights, self._security_by_ticker["GLD"], 0.08)
        self._add_weight(weights, self._security_by_ticker["SHY"], 0.04)
        return self._apply_cash_buffer(weights)

    def _weak_risk_on_weights(self):
        selected = self._select_top_growth_securities(2, 2)
        weights = {}
        if len(selected) >= 1:
            self._add_weight(weights, selected[0], 0.25)
        if len(selected) >= 2:
            self._add_weight(weights, selected[1], 0.15)
        self._add_weight(weights, self._security_by_ticker["GLD"], 0.30)
        self._add_weight(weights, self._security_by_ticker["SHY"], 0.20)
        self._add_weight(weights, self._security_by_ticker["XLV"], 0.10)
        return self._apply_cash_buffer(weights)

    def _risk_off_weights(self):
        return self._apply_cash_buffer({self._security_by_ticker["SHY"]: 0.50, self._security_by_ticker["GLD"]: 0.38, self._security_by_ticker["IEF"]: 0.12})

    def _select_top_growth_securities(self, count, max_high_beta):
        # Score each growth name by VaA momentum plus a trend bonus, normalized by volatility.
        score_by_security = {}
        for security in self._growth_securities:
            volatility = security.std_63.current.value
            if security.price <= 0 or volatility <= 0:
                score_by_security[security] = -100
                continue
            trend_bonus = 0
            if security.price > security.sma_50.current.value:
                trend_bonus += 0.03
            if security.price > security.sma_100.current.value:
                trend_bonus += 0.04
            if security.price > security.sma_200.current.value:
                trend_bonus += 0.06
            score_by_security[security] = (self._vaa_score(security) + trend_bonus) / max(0.03, volatility / security.price)
        # Take the highest scorers while capping how many high-beta names can be chosen.
        selected = []
        high_beta_count = 0
        for security in sorted(score_by_security, key=lambda s: score_by_security[s], reverse=True):
            if score_by_security[security] <= 0:
                continue
            if security.is_high_beta and high_beta_count >= max_high_beta:
                continue
            selected.append(security)
            if security.is_high_beta:
                high_beta_count += 1
            if len(selected) == count:
                break
        if len(selected) == 0:
            selected.append(self._spy)
        return selected

    def _vaa_score(self, security):
        return 12 * security.rocp_21.current.value + 4 * security.rocp_63.current.value + 2 * security.rocp_126.current.value + security.rocp_252.current.value

    def _add_weight(self, weights, security, weight):
        weights[security] = weights.get(security, 0) + weight

    def _apply_cash_buffer(self, weights):
        return {security: weight * (1 - 0.02) for security, weight in weights.items()}