from datetime import timedelta
from System.Drawing import Color

class OptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2019, 5, 1)
        self.SetCash(20000)
        self.syl = 'SPY'
        equity = self.AddEquity(self.syl, Resolution.Daily)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.macd = self.MACD(self.syl, 12, 26, 9, MovingAverageType.Exponential, Resolution.Daily)
        self.underlyingsymbol = equity.Symbol
        # use the underlying equity as the benchmark
        self.SetBenchmark(equity.Symbol)
    
        # set up a chart to display our buy and sell dates
        self.stockPlot = Chart('MACD')
        self.stockPlot.AddSeries(Series('Buy', SeriesType.Scatter, '$', Color.Green, ScatterMarkerSymbol.Triangle))
        self.stockPlot.AddSeries(Series('Close', SeriesType.Scatter, '$', Color.Red, ScatterMarkerSymbol.TriangleDown))
        self.AddChart(self.stockPlot)
        
        self.SetWarmUp(26)
        
        self.i = 0
        
        
    def OnData(self,slice):
        
        if self.IsWarmingUp: return
            
        # if there is a MACD long signal, buy a call     
        
        if self.macd.Current.Value > self.macd.Signal.Current.Value:
            if not self.i:
                self.BuyCall()
                self.Plot("MACD",'Buy',self.macd.Current.Value)
                self.i = 1
        else:
            if self.i:
                try:
                    self.MarketOrder(self.contract,-1)
                except:
                    pass
                
                self.Plot("MACD",'Close',self.macd.Current.Value)
                self.i = 0
            
        self.Plot("MACD",'MACD', self.macd.Current.Value)
        self.Plot("MACD",'signal',self.macd.Signal.Current.Value)
        
        try:
            self.Plot("Option",'price',self.contract.Value)
        except:
            pass
               
    def BuyCall(self):
        contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
        if len(contracts) == 0: return
        filtered_contracts = self.InitialFilter(self.underlyingsymbol, contracts, -3, 3, 0, 30)
        call = [x for x in filtered_contracts if x.ID.OptionRight == 0] 
        # sorted the contracts according to their expiration dates and choose the ATM options
        contracts = sorted(sorted(call, key = lambda x: abs(self.Securities[self.syl].Price - x.ID.StrikePrice)), 
                                        key = lambda x: x.ID.Date, reverse=True)
        self.contract = contracts[0]
        self.AddOptionContract(self.contract, Resolution.Minute)
        self.MarketOrder(self.contract, 1)
        
        

    def InitialFilter(self, underlyingsymbol, symbol_list, min_strike_rank, max_strike_rank, min_expiry, max_expiry):
        
        ''' This method is an initial filter of option contracts
            according to the range of strike price and the expiration date '''
            
        if len(symbol_list) == 0 : return
        # fitler the contracts based on the expiry range
        contract_list = [i for i in symbol_list if min_expiry < (i.ID.Date.date() - self.Time.date()).days < max_expiry]
        # find the strike price of ATM option
        atm_strike = sorted(contract_list,
                            key = lambda x: abs(x.ID.StrikePrice - self.Securities[underlyingsymbol].Price))[0].ID.StrikePrice
        strike_list = sorted(set([i.ID.StrikePrice for i in contract_list]))
        # find the index of ATM strike in the sorted strike list
        atm_strike_rank = strike_list.index(atm_strike)
        try: 
            min_strike = strike_list[atm_strike_rank + min_strike_rank]
            max_strike = strike_list[atm_strike_rank + max_strike_rank]
        except:
            min_strike = strike_list[0]
            max_strike = strike_list[-1]
           
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice >= min_strike and i.ID.StrikePrice <= max_strike]

        return filtered_contracts