# region imports
from AlgorithmImports import *
# endregion


class VIXAlphaModel(AlphaModel):

    _symbols = []

    def __init__(self, algorithm, lookback_days, long_percentile, short_percentile):
        self._vix = algorithm.add_data(CBOE, "VIX", Resolution.DAILY).symbol
        self._window = RollingWindow[float](lookback_days)
        self._long_percentile = long_percentile
        self._short_percentile = short_percentile

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        if data.contains_key(self._vix):
            # Get the current VIX value
            price = data[self._vix].price

            # Get the trailing VIX values
            self._window.add(price)
            if not self._window.is_ready:
                return []
            history_close = [i for i in self._window]

            # Check if the current VIX value is in an area of extreme relative to trailing values
            if price > np.percentile(history_close, self._long_percentile):
                direction = InsightDirection.UP
            elif price < np.percentile(history_close, self._short_percentile):
                direction = InsightDirection.DOWN
            else:
                return []
            # Emit insights if the current VIX value is extreme
            return [Insight.price(symbol, timedelta(365), direction) for symbol in self._symbols]
        return []

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            self._symbols.append(security.symbol)

# region imports
from AlgorithmImports import *

from alpha import VIXAlphaModel
# endregion


class VIXPredictsStockIndexReturns(QCAlgorithm):

    _undesired_symbols_from_previous_deployment = []
    _checked_symbols_from_previous_deployment = False
    _previous_direction_by_symbol = {}
    _week = -1

    def initialize(self):
        self.set_start_date(2023, 3, 1)
        self.set_end_date(2024, 3, 1)
        self.set_cash(1_000_000)

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.add_universe_selection(ManualUniverseSelectionModel([Symbol.create("OEF", SecurityType.EQUITY, Market.USA)]))

        lookback_days = self.get_parameter("lookback_days", 504)
        self.add_alpha(VIXAlphaModel(
            self, 
            lookback_days,
            self.get_parameter("long_percentile", 90),
            self.get_parameter("short_percentile", 10)
        ))

        self.settings.rebalance_portfolio_on_insight_changes = False
        self.settings.rebalance_portfolio_on_security_changes = False
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel(self._rebalance_func))

        self.add_risk_management(NullRiskManagementModel())
        
        self.set_execution(ImmediateExecutionModel())
        
        self.set_warm_up(timedelta(lookback_days*2))


    def _rebalance_func(self, time):
        # Rebalance when all of the following are true:
        #  - Not warming up
        #  - There is QuoteBar data in the current slice
        #  - It's a new week or the insight direction has changed
        direction_by_symbol = {insight.symbol: insight.direction for insight in self.insights}
        if (not self.is_warming_up and self.current_slice.quote_bars.count > 0 and 
            (direction_by_symbol != self._previous_direction_by_symbol or self.time.date().isocalendar()[1] != self._week)):
            self._week = self.time.date().isocalendar()[1]
            self._previous_direction_by_symbol = direction_by_symbol
            return time
        return None
        

    def on_data(self, data):
        # Exit positions that aren't backed by existing insights.
        # If you don't want this behavior, delete this method definition.
        if not self.is_warming_up and not self._checked_symbols_from_previous_deployment:
            for security_holding in self.portfolio.values():
                if not security_holding.invested:
                    continue
                symbol = security_holding.symbol
                if not self.insights.has_active_insights(symbol, self.utc_time):
                    self._undesired_symbols_from_previous_deployment.append(symbol)
            self._checked_symbols_from_previous_deployment = True
        
        for symbol in self._undesired_symbols_from_previous_deployment:
            if self.is_market_open(symbol):
                self.liquidate(symbol, tag="Holding from previous deployment that's no longer desired")
                self._undesired_symbols_from_previous_deployment.remove(symbol)