# region imports
from AlgorithmImports import *
from datetime import timedelta
from pandas.core.frame import DataFrame
# endregion

class OptionTradeFunction(Enum):
    BUY = 1
    SELL = 2

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

class Options(QCAlgorithm):

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

    _DTE: int = 35
    _OTM: float = 0.01
    _vix_rank_lookback: int = 150
    _percentage_traded: float = 1.0
    _options_alloc: int = 100   # options quantity = equity symbol quantity / _options_alloc
    
    _safe_asset_ticker: str = 'IEF'

    _traded_option_right: OptionRight = OptionRight.PUT
    _option_trade_function: OptionTradeFunction = OptionTradeFunction.SELL

    def initialize(self) -> None:
        self.init_cash: int = 100_000

        self.SetStartDate(2015, 1, 1)
        self.SetCash(self.init_cash)

        self._equity: Equity = self.add_equity("SPY", Resolution.MINUTE)
        self._equity.set_data_normalization_mode(DataNormalizationMode.RAW)
        self._benchmark_values: List[float] = []
        
        self._market: Symbol = self._equity.symbol
        self._trade_fn: Callable[Symbol, float] = self.buy if self._option_trade_function == OptionTradeFunction.BUY else self.sell
        self._move_to_safe_asset: bool = False

        self._contract: Symbol|str = str()
        self._contracts_added: Set = set()

        # safe asset
        self._safe_asset: Symbol = self.add_equity(self._safe_asset_ticker, Resolution.MINUTE).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]
            self.set_warm_up(self._vix_rank_lookback, Resolution.DAILY)
        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.settings.daily_precise_end_time = False

        self._recent_month: int = -1

        self.schedule.on(self.date_rules.every_day(self._market), \
                        self.time_rules.after_market_open(self._market, 30), \
                        self._VIX_filter)

    def _VIX_filter(self) -> None:
        if self._vix_filter_type == VixFilterType.VVIX:
            vvix: float = self.securities[self._signal_assets[0]].price
            if vvix > self._vix_filter_value_threshold:
                self._move_to_safe_asset = True
        elif self._vix_filter_type == VixFilterType.VIX_RANK:
            vix_rank: float = self._get_VIX_rank(self._signal_assets[0], self._vix_rank_lookback)
            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:
                self._move_to_safe_asset = True

    def on_end_of_day(self, symbol: Symbol) -> None:
        # print benchmark in main equity plot
        mkt_price_df: DataFrame = self.history(self._equity.symbol, 2, Resolution.DAILY)
        if not mkt_price_df.empty:
            benchmark_price: float = mkt_price_df['close'].unstack(level= 0).iloc[-1]
            if len(self._benchmark_values) == 2:
                self._benchmark_values[-1] = benchmark_price
                benchmark_perf: float = self.init_cash * (self._benchmark_values[-1] / self._benchmark_values[0])
                self.plot('Strategy Equity', self._equity.symbol, benchmark_perf)
            else:
                self._benchmark_values.append(benchmark_price)

    def on_data(self, slice: Slice) -> None:
        if self.is_warming_up:
            return
        
        if not self._move_to_safe_asset:
            # liquidate safe asset holdings
            if self.portfolio[self._safe_asset].invested:
                self.liquidate(self._safe_asset)

            # trade options with market
            self._trade_option(slice)

            if not self.portfolio[self._market].invested:
                self.set_holdings(self._market, self._percentage_traded)
        else:
            self._move_to_safe_asset = False
            # liquidate market holdings and options
            if self.portfolio[self._market].invested:
                self.liquidate(self._market)

            option_invested: List = [
                x.key for x in self.portfolio if x.value.invested and x.value.type == SecurityType.OPTION
            ]
            for symbol in option_invested:
                self.liquidate(symbol)

            # trade safe asset
            if not self.portfolio[self._safe_asset].invested:
                self.set_holdings(self._safe_asset, self._percentage_traded)

        # monthly option trade
        if self._recent_month != self.time.month:
            self._recent_month = self.time.month

            if self._contract:
                self.remove_option_contract(self._contract)
                self._contract = str()

    def _trade_option(self, slice: Slice) -> None:
        if self._contract == str():
            self._contract = self._filter_options(slice, self._traded_option_right)
            return

        elif not self.portfolio[self._contract].invested and slice.contains_key(self._contract):
            self._trade_fn(self._contract, round(self.portfolio[self._market].quantity / self._options_alloc))

    def _filter_options(self, slice: Slice, option_right: OptionRight) -> Symbol|str:
        contracts: List[Symbol] = self.option_chain_provider.get_option_contract_list(self._market, slice.time)
        
        underlying_price: float = self.securities[self._market].price
        
        if option_right == OptionRight.CALL:
            options: List[Symbol] = [
                i for i in contracts if i.id.option_right == OptionRight.CALL and   
                i.id.strike_price >= underlying_price * (1 + self._OTM) and
                i.id.date >= slice.time + timedelta(days=self._DTE)
            ]
        else:
            options: List[Symbol] = [
                i for i in contracts if i.id.option_right == OptionRight.PUT and   
                i.id.strike_price <= underlying_price * (1 - self._OTM) and
                i.id.date >= slice.time + timedelta(days=self._DTE)
            ]

        if len(options) > 0:
            expiry: datetime.datetime = min(list(map(lambda x: x.id.date, options)))

            index: int = 0 if option_right == OptionRight.PUT else -1
            contract: Symbol = sorted([o for o in options if o.id.date == expiry], key = lambda x: underlying_price - x.id.strike_price)[index]

            if contract not in self._contracts_added:
                self._contracts_added.add(contract)
                self.add_option_contract(contract, Resolution.MINUTE)

            return contract
        else:
            return str()

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