# region imports
from AlgorithmImports import *
# endregion


class HybridModel(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100000)
        self._qqq = self.add_equity("QQQ", Resolution.DAILY).symbol
        self._smh = self.add_equity("SMH", Resolution.DAILY).symbol
        self._gld = self.add_equity("GLD", Resolution.DAILY).symbol
        self._bil = self.add_equity("BIL", Resolution.DAILY).symbol
        self._risk_assets = [self._qqq, self._smh, self._gld]
        self._symbols = [self._qqq, self._smh, self._gld, self._bil]
        self._window_by_symbol = {symbol: RollingWindow[float](220) for symbol in self._symbols}
        self.set_warm_up(timedelta(days=365), Resolution.DAILY)
        self.schedule.on(self.date_rules.week_start(self._qqq), self.time_rules.at(8, 0), self._rebalance)

    def on_data(self, data):
        for symbol in self._symbols:
            if symbol in data.bars:
                self._window_by_symbol[symbol].add(float(data.bars[symbol].close))

    def on_warmup_finished(self):
        self._rebalance()

    def _rebalance(self):
        if self.is_warming_up:
            return
        # Score each risk asset by 6-month momentum divided by realized volatility.
        score_by_symbol = {}
        vol_by_symbol = {}
        for symbol in self._risk_assets:
            window = self._window_by_symbol[symbol]
            if not window.is_ready:
                continue
            price = window[0]
            old_price = window[126]
            if old_price <= 0 or price <= sum(window[i] for i in range(200)) / 200:
                continue
            momentum = price / old_price - 1.0
            if momentum <= 0:
                continue
            returns = [window[i] / window[i + 1] - 1.0 for i in range(63) if window[i + 1] > 0]
            if not returns:
                continue
            volatility = float(np.std(returns))
            if volatility > 0:
                score_by_symbol[symbol] = momentum / volatility
                vol_by_symbol[symbol] = volatility
        # Allocate 75% of the book across the top two names by inverse volatility.
        targets = {}
        remaining_weight = 0.98
        selected = sorted(score_by_symbol, key=lambda s: score_by_symbol[s])[-2:]
        inverse_volatility_sum = sum(1.0 / vol_by_symbol[symbol] for symbol in selected)
        if inverse_volatility_sum > 0:
            for symbol in selected:
                targets[symbol] = 0.75 * (1.0 / vol_by_symbol[symbol]) / inverse_volatility_sum
                remaining_weight -= targets[symbol]
        # Add a 20% boost to the deepest short-term pullback that is still above trend.
        pullback_symbol = None
        worst_pullback = 0.0
        for symbol in self._risk_assets:
            window = self._window_by_symbol[symbol]
            if not window.is_ready:
                continue
            price = window[0]
            five_day_price = window[5]
            if five_day_price <= 0 or price <= sum(window[i] for i in range(200)) / 200:
                continue
            pullback = price / five_day_price - 1.0
            if pullback < -0.02 and pullback < worst_pullback:
                pullback_symbol = symbol
                worst_pullback = pullback
        if pullback_symbol is not None:
            targets[pullback_symbol] = targets.get(pullback_symbol, 0) + 0.20
            remaining_weight -= 0.20
        if remaining_weight > 0:
            targets[self._bil] = remaining_weight
        for holding in list(self.portfolio.values()):
            if holding.invested and holding.symbol not in targets:
                self.set_holdings(holding.symbol, 0)
        for symbol, weight in targets.items():
            self.set_holdings(symbol, weight)