# region imports
from AlgorithmImports import *
from collections import deque
# endregion


class MaxEffectShortAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_cash(1_000_000)
        self.settings.seed_initial_prices = True

        self.universe_settings.resolution = Resolution.DAILY
        self.universe_settings.data_normalization_mode = DataNormalizationMode.ADJUSTED

        self._selection_data_by_symbol = {}
        self._universe = self.add_universe(self._select_assets)

        self._vix_symbol = self.add_index("VIX").symbol

        self.set_warm_up(21, Resolution.DAILY)

        self._current_targets = {}
        self._last_vix_value = None

        self.schedule.on(
            self.date_rules.month_start('SPY', 1),
            self.time_rules.at(8, 0),
            self._rebalance
        )

    def on_data(self, data: Slice) -> None:
        if not self.securities[self._vix_symbol].has_data or self._vix_symbol not in data:
            return

        vix_value = data[self._vix_symbol].value
        if self._last_vix_value is None:
            self._last_vix_value = vix_value
            return

        # Check if VIX crossed a threshold since the previous bar
        vix_crossed = (
            (self._last_vix_value < 15 and vix_value >= 15) or
            (self._last_vix_value >= 15 and vix_value < 15) or
            (self._last_vix_value < 30 and vix_value >= 30) or
            (self._last_vix_value >= 30 and vix_value < 30)
        )

        self._last_vix_value = vix_value

        if not vix_crossed:
            return

        if not self._current_targets:
            return

        new_exposure = self._get_target_exposure()
        value_by_symbol = {symbol: self.portfolio[symbol].holdings_value for symbol in self._current_targets if self.securities[symbol].price}
        total_holdings_value = sum(value_by_symbol.values())
        weight_by_symbol = {symbol: holdings_value / total_holdings_value for symbol, holdings_value in value_by_symbol.items()}
        scaled_targets = [PortfolioTarget(s, w * new_exposure) for s, w in weight_by_symbol.items()]
        self.set_holdings(scaled_targets, liquidate_existing_holdings=True)

    def _select_assets(self, fundamentals: List[Fundamental]) -> List[Symbol]:
        today_symbols = set()
        selected_symbols = set()

        for f in fundamentals:
            if not self._is_valid_universe(f):
                continue

            today_symbols.add(f.symbol)
            data = self._selection_data_by_symbol.setdefault(f.symbol, SelectionData())
            if data.update(f):
                selected_symbols.add(f.symbol)

        for symbol in list(self._selection_data_by_symbol.keys()):
            if symbol not in today_symbols:
                del self._selection_data_by_symbol[symbol]

        if self.is_warming_up:
            return []

        return list(selected_symbols)

    def _is_valid_universe(self, f: Fundamental) -> bool:
        if f.company_reference.is_reit:
            return False
        if f.security_reference.is_depositary_receipt:
            return False
        if not f.has_fundamental_data:
            return False
        if f.price < 5:
            return False
        if f.market_cap < 1_000_000_000:
            return False
        return True

    def _get_target_exposure(self) -> float:
        vix_security = self.securities.get(self._vix_symbol)
        if vix_security is None or not vix_security.has_data:
            return 1.0
        vix_value = vix_security.price
        if vix_value <= 15:
            return 1.5
        elif vix_value >= 30:
            return 1.0
        return 1.5 - (vix_value - 15) / 15.0 * 0.5

    def _rebalance(self) -> None:
        symbols = [s for s in self._universe.selected if s in self._selection_data_by_symbol]
        if not symbols:
            return

        max_5_by_symbol = {}
        for s in symbols:
            data = self._selection_data_by_symbol[s]
            if not data.is_ready:
                continue
            if data.dollar_volume_sma.current.value < 5_000_000:
                continue
            if not self.securities[s].has_data or not self.securities[s].price:
                continue
            max_5_by_symbol[s] = data.max_5

        if not max_5_by_symbol:
            return

        sorted_items = sorted(max_5_by_symbol.items(), key=lambda x: x[1])
        n = len(sorted_items)
        leg_size = min(25, n // 10)

        long_symbols = [s for s, _ in sorted_items[:leg_size]]

        targets = {}
        if long_symbols:
            exposure = self._get_target_exposure()
            long_weight = exposure / len(long_symbols)
            for s in long_symbols:
                targets[s] = long_weight

        self._current_targets = targets
        target_list = [PortfolioTarget(s, w) for s, w in targets.items()]
        self.set_holdings(target_list, liquidate_existing_holdings=True)


class SelectionData:
    def __init__(self) -> None:
        self._last_adjusted_price = None
        self._returns = deque(maxlen=21)
        self.dollar_volume_sma = SimpleMovingAverage(21)

    def update(self, f: Fundamental) -> bool:
        self.dollar_volume_sma.update(f.end_time, f.dollar_volume)

        if self._last_adjusted_price is not None and self._last_adjusted_price > 0:
            ret = f.adjusted_price / self._last_adjusted_price - 1.0
            self._returns.append(ret)

        self._last_adjusted_price = f.adjusted_price

        return self.is_ready

    @property
    def is_ready(self) -> bool:
        return len(self._returns) == 21 and self.dollar_volume_sma.is_ready

    @property
    def max_5(self) -> float:
        returns = sorted(self._returns, reverse=True)
        return float(sum(returns[:5]) / 5.0)