from datetime import timedelta

class OptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 01, 01)
        self.SetEndDate(2017, 01, 15)
        self.SetCash(100000)
        equity = self.AddEquity("VXX", Resolution.Minute)
        self.symbol = equity.Symbol
        self.last_slice = None

        # warm up for 121 days
        self.SetWarmUp(timedelta(121))
        
        # filter the option contracts we will trade
        option = self.AddOption(str(self.symbol))
        # set our strike/expiry filter for this option chain
        option.SetFilter(self.UniverseFunc)
        #option.SetFilter(-2, +20, TimeSpan.FromDays(5),TimeSpan.FromDays(9))
     
        # run rebalance every day before market close
        self.Schedule.On(self.DateRules.Every([DayOfWeek.Friday]), self.TimeRules.BeforeMarketClose(str(self.symbol), 2), Action(self.rebalance))

    
    def OnData(self,slice):
        if self.IsWarmingUp: return
        # if self.Time.weekday() != 4 or self.Time.hour != 12 or self.Time.minute != 58: return
        
        self.last_slice = slice

    
    def rebalance(self):
        self.Debug("rebalance being called")
        
        # sell everything first
        self.Liquidate()
        
        for i in self.last_slice.OptionChains:
            chain = i.Value

            # differentiate the call and put options
            calls = [x for x in chain if x.Right == 0]
            puts = [x for x in chain if x.Right == 1]

            # get all OTM calls
            otm_calls = [x for x in calls if x.UnderlyingLastPrice - x.Strike < 0]
            if len(otm_calls) == 0: return
            # get closest ATM call
            atm_call = sorted(otm_calls, key = lambda x: x.UnderlyingLastPrice - x.Strike)[0]
            
            # get OTM call at least 10% OTM
            fotm_calls = [x for x in otm_calls if float(x.UnderlyingLastPrice)*float(1.15) - float(x.Strike) < 0]
            fotm_call = sorted(fotm_calls, key = lambda x: x.UnderlyingLastPrice - x.Strike)[0]

            # get all OTM puts
            otm_puts = [x for x in puts if x.UnderlyingLastPrice - x.Strike > 0]
            
            # get 2nd closest ATM put
            put = sorted(otm_puts, key = lambda x: x.Strike, reverse = True)[1]

            # sell 80% ATM calls
            atm_call_amount = float(self.Portfolio.TotalPortfolioValue) * -0.8
            atm_call_contracts = atm_call_amount / (float(atm_call.UnderlyingLastPrice)*100.0)
            self.MarketOrder(atm_call.Symbol, atm_call_contracts)

            # buy 80% Far out of the money calls
            fotm_call_amount = float(self.Portfolio.TotalPortfolioValue) * 0.8
            fotm_call_contracts = fotm_call_amount / (float(fotm_call.UnderlyingLastPrice)*100.0)
            self.MarketOrder(fotm_call.Symbol, fotm_call_contracts)

            # buy 32% puts
            put_amount = float(self.Portfolio.TotalPortfolioValue) * 0.32
            put_contracts = put_amount / (float(put.UnderlyingLastPrice)*100.0)
            self.MarketOrder(put.Symbol, put_contracts)
     
    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys().Strikes(-2, 20).Expiration(TimeSpan.FromDays(5),TimeSpan.FromDays(9))