# region imports
from AlgorithmImports import *
# endregion

class Magnificent7(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2020, 1, 10)
        self.set_end_date(2020, 1, 30) # Set End Date
        self.set_cash(1000000)
        self.underlying_tickers = ['GOOG', 'AAPL', 'AMZN', 'META', 'MSFT', 'NVDA', 'TSLA']
        
        # Initialize variables
        self.option_symbol = {} # Dictionary to store symbols
        self.sma_dict = {} # Dictionary to store the SMA for each symbol
        self.current_IV_dict = {}
        self.IVR_52W_low_dict = {}
        self.IVR_52W_high_dict = {}
        self.min_IV_52W_low = {}
        self.max_IV_52W_high = {}
                
        self.SetBenchmark("SPY")
        # Set window for simple moving average
        self.sma_window = 100        

        for ticker in self.underlying_tickers:
            # Add equity to universe
            equity = self.add_equity(ticker)
            # Add equity option contracts to universe and create list
            option = self.add_option(ticker)
            # Filters the option chain so that only 
            # strikes that are +$100 above current underlying price and
            # strikes $0 below current underlying price option and 
            # within 0 to 45 days to expiration (DTE) of current date
            # are available. (Speeds up code!)
            option.set_filter(0, 100, timedelta(0), timedelta(45))
            self.option_symbol[ticker] = option.symbol
            # Add SMA indicator to dictionary based on ticker symbol
            self.sma_dict[ticker] = self.sma(equity.Symbol, self.sma_window)
            # Dictionaries to keep track of 52 week high and low IV and current IV
            self.IVR_52W_low_dict[ticker] = []
            self.IVR_52W_high_dict[ticker] = []
            self.current_IV_dict[ticker] = []
            # Set initial values for max and min values of IV
            self.min_IV_52W_low[ticker] = 0
            self.max_IV_52W_high[ticker] = 0     

        # Run function every day after market closes to record daily high and low IV        
        self.schedule.on(self.date_rules.every_day("GOOG"),
                 self.time_rules.before_market_close("GOOG", -1),
                 self.after_market_close)
            
    def on_data(self, slice: Slice) -> None:
        for ticker in self.underlying_tickers:
            if not self.is_market_open(self.option_symbol[ticker]): return
            bar = slice.bars.get(ticker)
            chain = slice.option_chains.get_value(self.option_symbol[ticker])
            if chain is None: continue # two ways to check
            # Sort the contracts to find at the money (ATM) contract with closest expiration
            contracts = sorted(sorted(sorted(chain, \
                            key = lambda x: abs(chain.underlying.price - x.strike)), \
                            key = lambda x: x.expiry, reverse=False), \
                            key = lambda x: x.right, reverse=True)
            # if found, trade it
            if len(contracts) == 0: return
            current_IV = contracts[0].ImpliedVolatility
            if not self.current_IV_dict[ticker]:
                self.current_IV_dict[ticker] = [current_IV]*2
            else:            
                current_IV_low = self.current_IV_dict[ticker][0]
                current_IV_high = self.current_IV_dict[ticker][1]
                if (current_IV < current_IV_low):
                    self.current_IV_dict[ticker][0] = current_IV                
                elif (current_IV > current_IV_high):
                    self.current_IV_dict[ticker][1] = current_IV
            if self.max_IV_52W_high[ticker] and self.min_IV_52W_low[ticker]:               
                IVR = (current_IV-self.min_IV_52W_low[ticker]) / (self.max_IV_52W_high[ticker]-self.min_IV_52W_low[ticker]) * 100
            else:
                IVR = None

            if bar and IVR:
                if not self.portfolio.invested and (bar.Close < self.sma_dict[ticker].Current.Value) and (IVR > 30):
                    self.log(f"Current IVR for {ticker}: {IVR}. Buy 100 shares and sell covered call.")
                    chain = slice.option_chains.get(self.option_symbol[ticker])
                    if chain is None: return
                    # we sort the contracts to identify call options
                    call = [x for x in chain if x.Right == OptionRight.Call]
                    # we sort the contracts by closest expiration
                    contracts = sorted(call, key = lambda x: x.expiry, reverse=False)
                    # we sort the contracts based on given delta
                    delta_contracts = sorted(contracts, key = lambda x: abs(x.Greeks.Delta - 0.25))
                    # if found, trade it
                    if len(delta_contracts) == 0: return
                    self.market_order(ticker, 100)
                    self.market_order(delta_contracts[0].symbol,-1)
                    self.market_on_close_order(ticker, -100)
                    self.market_on_close_order(delta_contracts[0].symbol, 1)


    def after_market_close(self) -> None:
        # self.log(f"Fired at: {self.time}")        
        for ticker in self.underlying_tickers:
            if not self.current_IV_dict[ticker]: continue
            if len(self.IVR_52W_high_dict[ticker]) > 52*5:
                self.IVR_52W_high_dict[ticker].pop(0)    
            if len(self.IVR_52W_low_dict[ticker]) > 52*5:
                self.IVR_52W_low_dict[ticker].pop(0)
            self.IVR_52W_low_dict[ticker].append(self.current_IV_dict[ticker][0])
            self.IVR_52W_high_dict[ticker].append(self.current_IV_dict[ticker][1])            
            self.max_IV_52W_high[ticker] = max(self.IVR_52W_high_dict[ticker])
            self.min_IV_52W_low[ticker] = min(self.IVR_52W_low_dict[ticker])
            if ticker == self.underlying_tickers[0]:
                current_avg_IV = 0.5 * (self.current_IV_dict[ticker][0]+self.current_IV_dict[ticker][1])
                avg_IVR = (current_avg_IV-self.min_IV_52W_low[ticker]) / (self.max_IV_52W_high[ticker]-self.min_IV_52W_low[ticker]) * 100            
                self.plot("IVR", "Avg IVR", avg_IVR)
            self.current_IV_dict[ticker]=[]