# region imports
from AlgorithmImports import *
# endregion
# Watch my Tutorial: https://youtu.be/Lq-Ri7YU5fU


class OptionChainProviderPutProtection(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.settings.seed_initial_prices = True
        self.settings.automatic_indicator_warm_up = True
        # parameters ------------------------------------------------------------
        self._days_before_exp = 2 # number of days before expiry to exit
        self._dte = 25 # target days till expiration
        self._otm = 0.01 # target percentage OTM of put
        self._lookback_iv = 150 # lookback length of IV indicator
        self._iv_lvl = 0.5 # enter position at this lvl of IV indicator
        self._percentage = 0.9 # percentage of portfolio for underlying asset
        self._options_alloc = 90 # 1 option for X num of shares (balanced would be 100)
        # ------------------------------------------------------------------------
        # Add the underlying asset.
        self._equity = self.add_equity("SPY", data_normalization_mode=DataNormalizationMode.RAW)
        # Add VIX data.
        self._vix = self.add_data(CBOE, "VIX")
        # Initialize the IV rank indicator.
        self._rank = 0
        # Create a member to trade the Option contract.
        self._contract = None
        # Schedule plotting 30 minutes after every market open.
        date_rule = self.date_rules.every_day(self._equity)
        time_rule = self.time_rules.after_market_open(self._equity, 30)
        self.schedule.on(date_rule, time_rule, self._plotting)
        # Update the IV Rank 30 minutes after every market open.
        self.schedule.on(date_rule, time_rule, self._vix_rank)
        # warmup for IV indicator of data
        self.set_warm_up(timedelta(self._lookback_iv)) 

    def _vix_rank(self):
        history = self.history(CBOE, self._vix, self._lookback_iv, Resolution.DAILY)
        # (Current - Min) / (Max - Min)
        self._rank = ((self._vix.price - min(history["low"])) / (max(history["high"]) - min(history["low"])))
 
    def on_data(self, data):
        if self.is_warming_up:
            return
        # Buy the underlying asset if we don't already hold it.
        if not self._equity.invested:
            self.set_holdings(self._equity, self._percentage)
        # Buy a put if VIX is relatively high.
        if self._rank > self._iv_lvl:
            self._buy_put(data)
        # Close the put before it expires.
        if (self._contract and 
            (self._contract.expiry - self.time) <= timedelta(self._days_before_exp)):
            self.liquidate(self._contract)
            self._contract = None

    def _buy_put(self, data):
        # Get the Option contract.
        if not self._contract:
            self._contract = self._options_filter(data)        
        # If not invested and we have Option data, buy the Option.
        elif not self.portfolio[self._contract].invested and self._contract in data:
            self.buy(self._contract, round(self._equity.holdings.quantity / self._options_alloc))

    def _options_filter(self, data):
        chain = self.option_chain(self._equity)
        # Filter the OTM put Options from the contract list,
        # which expire close to self._dte num of days from now.
        otm_puts = [
            c for c in chain 
            if (c.right == OptionRight.PUT and
                self._equity.price - c.strike > self._otm * self._equity.price and
                self._dte - 8 < (c.expiry - data.time).days < self._dte + 8)
        ]
        if otm_puts:
            # Sort the contracts by closest to self._dte days 
            # from now and desired strike, and pick the first one.
            contract = sorted(
                sorted(otm_puts, key=lambda x: abs((x.expiry - self.time).days - self._dte)),
                key=lambda x: self._equity.price - x.strike
            )[0]
            self.add_option_contract(contract)
            return contract

    def _plotting(self):
        # plot IV indicator
        self.plot("Vol Chart", "Rank", self._rank)
        # plot indicator entry level
        self.plot("Vol Chart", "lvl", self._iv_lvl)
        # plot underlying's price
        self.plot("Data Chart", str(self._equity.symbol), self._equity.close)
        # plot strike of put option
        option_invested = [
            symbol for symbol, holding in self.portfolio.items() 
            if holding.invested and holding.type==SecurityType.OPTION
        ]
        if option_invested:
            self.plot("Data Chart", "strike", option_invested[0].id.strike_price)