from datetime import timedelta
from QuantConnect.Securities.Option import OptionPriceModels

class LongStrangleAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 4, 1)
        self.SetEndDate(2017, 5, 30)
        self.SetCash(100000)
        
        self.SetSecurityInitializer(self.SecurityInitializer)
        
        equity = self.AddEquity("GOOG", Resolution.Minute)
        option = self.AddOption("GOOG", Resolution.Minute)
        option.PriceModel  =  OptionPriceModels.CrankNicolsonFD() 
      
        equity.FeeModel = ConstantFeeModel(10.65) 
        #self.SetBrokerageModel(InteractiveBrokersBrokerageModel())
        self.symbol = option.Symbol
        # set our strike/expiry filter for this option chain
        option.SetFilter(-15, 15, timedelta(30), timedelta(60))
        # use the underlying equity GOOG as the benchmark
        self.SetBenchmark(equity.Symbol)
    
    def SecurityInitializer(self, security):
        if security.Type == SecurityType.Option:
            security.FeeModel = ConstantFeeModel(10.65)
        elif security.Type == SecurityType.Equity:
            #...
            pass
        
    def OnData(self,slice):
        optionchain = slice.OptionChains
        for i in slice.OptionChains:
            if i.Key != self.symbol: continue
            chains = i.Value
            contract_list = [x for x in chains]
            # if there is no contracts in this optionchain, pass the instance
            if (slice.OptionChains.Count == 0) or (len(contract_list) == 0): return   
             # if there is no securities in portfolio, trade the options 
            if not self.Portfolio.Invested: 
                self.TradeOptions(optionchain)

    def TradeOptions(self,optionchain):
    
        for i in optionchain:
            if i.Key != self.symbol: continue
            chain = i.Value
            # sorted the optionchain by expiration date and choose the furthest date
            expiry = sorted(chain,key = lambda x: x.Expiry, reverse=True)[0].Expiry
            # filter the call options from the contracts expires on that date
            call = [i for i in chain if i.Expiry == expiry and i.Right == 0]
            # sorted the contracts according to their strike prices 
            call_contracts = sorted(call,key = lambda x: x.Strike)    
            if len(call_contracts) == 0: continue
            # choose the deep OTM call option
            self.call = call_contracts[-1]
            # select the put options which have the same expiration date with the call option 
            # sort the put options by strike price
            put_contracts = sorted([i for i in chain if i.Expiry == expiry and i.Right == 1], key = lambda x: x.Strike)
            # choose the deep OTM put option
            self.put = put_contracts[0]
                    
            self.Buy(self.call.Symbol ,1)
            self.Buy(self.put.Symbol ,1)
    
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))