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


class MaxEffectAlgorithm(QCAlgorithm):

    def initialize(self) -> None:
        self.set_cash(1_000_000)
        self.set_start_date(self.end_date - timedelta(5*365))
        self.settings.seed_initial_prices = True
        # Add a universe of US Equities.
        self._selection_data_by_symbol = {}
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self._select_assets)
        # Add VIX to drive portfolio leverage decisions.
        self._vix = self.add_data(CBOE, 'VIX', Resolution.DAILY)
        self._low_vix = 15
        self._high_vix = 30
        self._previous_vix = None
        self._current_targets = []
        # Define some parameters.
        self._top_returns = self.get_parameter('top_returns', 5)
        self._max_portfolio_size = self.get_parameter('max_portfolio_size', 25)
        # Add a warm-up period so the algorithm trades right away.
        self.set_warm_up(timedelta(45))

    def on_warmup_finished(self) -> None:
        # Add a Scheduled Event to rebalance at the start of each month.
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self.date_rules.month_start('SPY', 1), time_rule, self._rebalance)
        # Rebalance today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def _select_assets(self, fundamentals: List[Fundamental]) -> List[Symbol]:
        # Update SelectionData objects.
        today_symbols = set()
        selected_symbols = set()
        for f in fundamentals:
            # Apply universe filters.
            if (f.company_reference.is_reit or
                f.security_reference.is_depositary_receipt or
                not f.has_fundamental_data or
                f.price < 5 or
                f.market_cap < 1_000_000_000):
                continue
            today_symbols.add(f.symbol)
            data = self._selection_data_by_symbol.setdefault(f.symbol, SelectionData(top_returns=self._top_returns))
            if data.update(f):
                selected_symbols.add(f.symbol)
        # Remove SelectionData objects for assets that have fallen out of the universe.
        for symbol in list(self._selection_data_by_symbol.keys()):
            if symbol not in today_symbols:
                del self._selection_data_by_symbol[symbol]
        return list(selected_symbols)

    def _rebalance(self) -> None:
        if self.is_warming_up:
            return
        securities = [self.securities[symbol] for symbol in self._universe.selected]
        if not securities:
            return
        # Get the factor of each asset.
        factor_by_security = {}
        for security in securities:
            data = self._selection_data_by_symbol[security.symbol]
            if security.price and data.dollar_volume_sma.current.value >= 5_000_000:
                factor_by_security[security] = data.factor
        # Select the subset of stocks in the bottom deciles of factor values.
        sorted_by_factor = sorted(factor_by_security, key=lambda security: factor_by_security[security])
        selection_size = min(self._max_portfolio_size, len(sorted_by_factor) // 10)
        securities = sorted_by_factor[:selection_size]
        self._current_securities = securities
        # Place trades to rebalance the portfolio.
        weight = 1 / len(securities)
        self._trade({security: weight for security in securities})

    def on_data(self, data: Slice) -> None:
        # Get the current VIX value.
        if self._vix not in data:
            return
        vix_value = data[self._vix].value
        if self._previous_vix is None:
            self._previous_vix = vix_value
            return
        if self.is_warming_up:
            return
        # Check if VIX crossed a threshold.
        vix_crossed = (
            (self._previous_vix < self._low_vix and vix_value >= self._low_vix) or
            (self._previous_vix >= self._low_vix and vix_value < self._low_vix) or
            (self._previous_vix < self._high_vix and vix_value >= self._high_vix) or
            (self._previous_vix >= self._high_vix and vix_value < self._high_vix)
        )
        self._previous_vix = vix_value
        if not (vix_crossed and self._current_securities):
            return
        # Adjust the portfolio leverage, while keeping relative weights between assets.
        value_by_security = {security: security.holdings.holdings_value for security in self._current_securities if security.price}
        total_holdings_value = sum(value_by_security.values())
        self._trade({security: holdings_value / total_holdings_value for security, holdings_value in value_by_security.items()})
    
    def _trade(self, weight_by_security):
        # Determine the portfolio leverage based on the VIX.
        vix = self._vix.price
        if not vix or vix >= self._high_vix:
            leverage = 1.0
        elif vix <= self._low_vix:
            leverage = 1.5
        else:
            leverage = 1.5 - (vix - self._low_vix) / self._low_vix * 0.5
        # Scale portfolio weights based on target portfolio exposure.
        targets = [PortfolioTarget(s, leverage * w ) for s, w in weight_by_security.items()]
        self.set_holdings(targets, liquidate_existing_holdings=True)


class SelectionData:

    def __init__(self, period=21, top_returns=5) -> None:
        self.dollar_volume_sma = SimpleMovingAverage(period)
        self._daily_return = RateOfChange(1)
        self._daily_return.window.size = period
        self._top_returns = top_returns

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

    @property
    def is_ready(self) -> bool:
        return self._daily_return.samples > self._daily_return.window.size

    @property
    def factor(self) -> float:
        # Get the mean of the top n trailing daily returns.
        return float(np.mean(sorted([x.value for x in self._daily_return.window])[-self._top_returns:]))