# region imports
from AlgorithmImports import *
from dateutil.relativedelta import relativedelta
# endregion

class VixFilterType(Enum):
    VVIX = 1
    VIX_RANK = 2
    VIX_RATIO = 3

class RecessionSignal(QCAlgorithm):

    _vix_filter_type: VixFilterType = VixFilterType.VIX_RANK
    _vix_filter_value_threshold: float = 0.5

    _safe_asset_ticker: str = 'IEF'
    _sma_period: int = 10 * 21

    def initialize(self) -> None:
        self.set_start_date(2000, 1, 1)
        self.set_cash(100_000)

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.daily_precise_end_time = False

        self._market: Symbol = self.add_equity("SPY", Resolution.MINUTE).symbol
        self._sma: SimpleMovingAverage = self.SMA(self._market, self._sma_period, Resolution.DAILY)

        # safe asset
        self._safe_asset: Symbol = self.add_equity(self._safe_asset_ticker, Resolution.MINUTE).symbol

        self.set_warm_up(self._sma_period, Resolution.DAILY)
        
        self._indpro: Symbol = self.add_data(YOYData, 'INDPRO_YOY', Resolution.DAILY).symbol
        self._rrsfs: Symbol = self.add_data(YOYData, 'RRSFS_YOY', Resolution.DAILY).symbol

        # subscribe to VIX filter asset/s
        if self._vix_filter_type == VixFilterType.VVIX:
            iv: Symbol = self.add_data(CBOE, "VVIX", Resolution.DAILY).symbol
            self._signal_assets = [iv]
        elif self._vix_filter_type == VixFilterType.VIX_RANK:
            iv: Symbol = self.add_data(CBOE, "VIX", Resolution.DAILY).symbol
            self._signal_assets = [iv]
        elif self._vix_filter_type == VixFilterType.VIX_RATIO:
            iv: Symbol = self.add_data(CBOE, "VIX", Resolution.DAILY).symbol
            iv_3m: Symbol = self.add_data(CBOE, "VIX3M", Resolution.DAILY).symbol
            self._signal_assets = [iv, iv_3m]

        self._rebalance_flag: bool = False
        self.schedule.on(self.date_rules.month_end(self._market), self.time_rules.before_market_close(self._market, 1), self._selection)
    
    def on_data(self, slice: Slice) -> None:
        if not self._rebalance_flag:
            return
        self._rebalance_flag = False

        if self.is_warming_up: return
        if not self._sma.is_ready: return
        if not all(slice.contains_key(symbol) for symbol in [self._market, self._safe_asset]): return
        if not all(
            self.securities.contains_key(symbol) and self.securities[symbol].get_last_data() for symbol in self._signal_assets
        ): return
        
        # trade safe asset
        move_to_safe_asset = False
        traded_asset: Symbol = self._market

        if self._vix_filter_type == VixFilterType.VVIX:
            vvix: float = self.securities[self._signal_assets[0]].price
            if vvix > self._vix_filter_value_threshold:
                move_to_safe_asset = True
        elif self._vix_filter_type == VixFilterType.VIX_RANK:
            vix_rank: float = self._get_VIX_rank(self._signal_assets[0])
            if vix_rank > self._vix_filter_value_threshold:
                move_to_safe_asset = True
        elif self._vix_filter_type == VixFilterType.VIX_RATIO:
            vix_ratio: float = self.securities[self._signal_assets[0]].price / self.securities[self._signal_assets[1]].price
            if vix_ratio > self._vix_filter_value_threshold:
                move_to_safe_asset = True

        if move_to_safe_asset:
            traded_asset = self._safe_asset
        else:
            # check end of custom datasets
            last_update_date: datetime.date = YOYData.get_last_update_date()
            if not all(self.securities[symbol].get_last_data() and symbol.value in last_update_date and self.time.date() < last_update_date[symbol.value] \
                for symbol in [self._indpro, self._rrsfs]):
                # signal no recession in case there's an end of custom data
                pass
            else:
                # signal recession
                if any(self.securities[symbol].price < 0. for symbol in [self._indpro, self._rrsfs]):
                    if slice[self._market].value <= self._sma.current.value:
                        traded_asset = self._safe_asset
        
        if not self.portfolio[traded_asset].invested:
            self.set_holdings(traded_asset, 1., True)

    def _selection(self) -> None:
        self._rebalance_flag = True

    def _get_VIX_rank(self, vix: Symbol, lookback: int = 150) -> float:
        history: DataFrame = self.history(CBOE, vix, lookback, Resolution.DAILY)
        rank: float = ((self.securities[vix].price - min(history["low"])) / (max(history["high"]) - min(history["low"])))
        return rank

class YOYData(PythonData):
    _last_update_date: Dict[str, datetime.date] = {}

    def GetSource(self, config: SubscriptionDataConfig, date: datetime, isLiveMode: bool) -> SubscriptionDataSource:
        return SubscriptionDataSource(f'data.quantpedia.com/backtesting_data/economic/{config.symbol.value}.csv', SubscriptionTransportMedium.REMOTE_FILE, FileFormat.CSV)

    @staticmethod
    def get_last_update_date() -> Dict[Symbol, datetime.date]:
       return YOYData._last_update_date

    def Reader(self, config: SubscriptionDataConfig, line: str, date: datetime, is_live_mode: bool) -> BaseData:
        data: YOYData = YOYData()
        data.symbol = config.symbol

        if not line[0].isdigit(): return None
        split = line.split(';')
        
        # Parse the CSV file's columns into the custom data class
        data.time = datetime.strptime(split[0], "%Y-%m-%d") + relativedelta(months=1)
        data.value = float(split[1])

        if config.symbol.value not in YOYData._last_update_date:
            YOYData._last_update_date[config.symbol.value] = datetime(1,1,1).date()

        if data.time.date() > YOYData._last_update_date[config.symbol.value]:
            YOYData._last_update_date[config.symbol.value] = data.time.date()
        
        return data