from AlgorithmImports import *
from datetime import timedelta


class WeeklyVaaMomentumScoreRotation(QCAlgorithm):
    def initialize(self) -> None:
        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._target_total_weight = 0.95
        self._risk_tickers = ["QQQ", "XLK", "XLE", "DBC", "GLD"]
        self._defensive_tickers = ["GLD", "BIL"]
        self._momentum_periods = [21, 63, 126, 252]
        self._momentum_weights = [12.0, 4.0, 2.0, 1.0]

        self._tickers = ["SPY", "QQQ", "XLK", "XLE", "DBC", "GLD", "BIL"]
        self._symbols_by_ticker = {}
        self._momentum_by_ticker = {}

        for ticker in self._tickers:
            symbol = self.add_equity(ticker, Resolution.DAILY).symbol
            self._symbols_by_ticker[ticker] = symbol
            self._momentum_by_ticker[ticker] = []
            for period in self._momentum_periods:
                self._momentum_by_ticker[ticker].append(self.rocp(symbol, period))

        self.set_warm_up(max(self._momentum_periods) + 20, Resolution.DAILY)

        self.schedule.on(
            self.date_rules.week_start(self._symbols_by_ticker["SPY"]),
            self.time_rules.at(8, 0),
            self._rebalance,
        )

    def on_warmup_finished(self) -> None:
        self._rebalance()

    def _rebalance(self) -> None:
        if self.is_warming_up:
            return

        if not self._indicators_are_ready():
            self.set_holdings(self._symbols_by_ticker["BIL"], self._target_total_weight)
            return

        selected = self._select_tickers()
        weight = self._target_total_weight / len(selected)

        targets = []
        for ticker in selected:
            symbol = self._symbols_by_ticker[ticker]
            targets.append(PortfolioTarget(symbol, weight))

        self.set_holdings(targets, liquidate_existing_holdings=True)

    def _indicators_are_ready(self) -> bool:
        for ticker in self._tickers:
            for indicator in self._momentum_by_ticker[ticker]:
                if not indicator.is_ready:
                    return False

        return True

    def _select_tickers(self) -> list[str]:
        if self._risk_regime_is_on():
            selected = self._rank_tickers(self._risk_tickers, 2, require_positive=True)
        else:
            selected = self._rank_tickers(self._defensive_tickers, 1, require_positive=True)

        if len(selected) == 0:
            selected = self._rank_tickers(self._defensive_tickers, 1, require_positive=False)

        if len(selected) == 0:
            return ["BIL"]

        return selected

    def _risk_regime_is_on(self) -> bool:
        return self._momentum_score("SPY") > 0

    def _rank_tickers(self, tickers: list[str], count: int, require_positive: bool) -> list[str]:
        ranked = []
        for ticker in tickers:
            score = self._momentum_score(ticker)
            if score > 0 or not require_positive:
                ranked.append((ticker, score))

        ranked.sort(key=lambda item: item[1], reverse=True)

        selected = []
        for item in ranked[:count]:
            selected.append(item[0])

        return selected

    def _momentum_score(self, ticker: str) -> float:
        score = 0.0
        indicators = self._momentum_by_ticker[ticker]
        for index in range(len(indicators)):
            score += self._momentum_weights[index] * indicators[index].current.value

        return score