from datetime import timedelta

class IronCondorAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 2, 1)
        self.SetEndDate(2017, 3, 1)
        self.SetCash(500000)
        equity = self.AddEquity("SPY", Resolution.Minute)
        option = self.AddOption("SPY", Resolution.Minute)
        self.symbol = option.Symbol
        option.SetFilter(self.UniverseFunc)
        # use the underlying equity GOOG as the benchmark
        self.SetBenchmark(equity.Symbol)
        
    def OnData(self,slice):
        self.TradeOptions(slice)


    def TradeOptions(self,slice):
        # If there is undelying assets in portfolio at expiration, liquidate the stocks in order to roll into new contracts
        if self.Portfolio["SPY"].Quantity != 0:
                self.Liquidate()
        
        if not self.Portfolio.Invested and self.Time.hour != 0 and self.Time.minute != 0: 
            for i in slice.OptionChains:
                chain = i.Value
                contract_list = [x for x in chain]
                # if there is no optionchain or no contracts in this optionchain, pass the instance
                if (slice.OptionChains.Count == 0) or (len(contract_list) == 0): 
                    return   
    
                # sorted the optionchain by expiration date and choose the furthest date
                expiry = sorted(chain,key = lambda x: x.Expiry)[-1].Expiry
                # filter the call and put options from the contracts
                call = [i for i in chain if i.Expiry == expiry and i.Right == 0]
                put = [i for i in chain if i.Expiry == expiry and i.Right == 1]
    
                # sorted the contracts according to their strike prices 
                call_contracts = sorted(call,key = lambda x: x.Strike)    
                put_contracts = sorted(put,key = lambda x: x.Strike)    
                if len(call_contracts) == 0 or len(put_contracts) == 0 : continue
    
                otm_put_lower = put_contracts[0]
                otm_put = put_contracts[10]
                otm_call = call_contracts[-10]
                otm_call_higher = call_contracts[-1]
                self.trade_contracts = [otm_call.Symbol,otm_call_higher.Symbol,otm_put.Symbol,otm_put_lower.Symbol]
    
                # if there is no securities in portfolio, trade the options 
                self.Buy(otm_put_lower.Symbol ,1)
                self.Sell(otm_put.Symbol ,1)
                self.Sell(otm_call.Symbol ,1)
                self.Buy(otm_call_higher.Symbol ,1)
    
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))

    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys().Strikes(-15, 15).Expiration(timedelta(0), timedelta(40))