from datetime import timedelta

class OptionsTest(QCAlgorithm):

    def Initialize(self):
        self.lastDayTraded = None
        
        self.SetStartDate(2020, 12, 2)
        self.SetEndDate(2020, 12, 6)
        self.SetCash(100000)

        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        option = self.AddOption("SPY")
        self.option_symbol = option.Symbol

        # set our strike/expiry filter for this option chain
        # FIXME: If I change timedelta(1) to timedelta(0), sometimes it grabs contracts that already expired
        option.SetFilter(lambda universe: universe.IncludeWeeklys().Strikes(-5, +5).Expiration(timedelta(1), timedelta(6)))

        # use the underlying equity as the benchmark
        self.SetBenchmark("SPY")
        
        # Rolling window of close prices for detecting dips
        self.close = RollingWindow[float](3)

    def OnData(self,slice):
        # Conditions when we don't want to open a new position
        if self.lastDayTraded == self.Time.date():
            return
        if not slice.ContainsKey("SPY"):
            return
        if self.Time.hour != 9 or self.Time.minute != 55:
            return

        for kvp in slice.OptionChains:
            if kvp.Key != self.option_symbol: continue
            chain = kvp.Value

            # we sort the call contracts to find at the money (ATM) contract with soonest 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=False) # Order calls first

            # If found, trade it
            if len(contracts) < 1:
                continue
            contract = contracts[0]
            
            # FIXME: Why are the prices so different (sometimes) here vs Yahoo Finance?
            self.Log(f"contract={contract}; bid={contract.BidPrice}; ask={contract.AskPrice}")
            
            # Open position and stop loss order
            size = 56
            symbol = contract.Symbol
            stopPrice = contract.AskPrice * 0.5
            self.MarketOrder(symbol, size)
            self.StopMarketOrder(symbol, -size, stopPrice)
            
            # Keep track if we traded today so we don't repeat
            self.lastDayTraded = self.Time.date()
            
            # Close position in 13 mins, cancels stop loss order if it's open
            self.Schedule.On(self.DateRules.Today,
                            self.TimeRules.At(self.Time.hour, self.Time.minute + 13),
                            self.Liquidate)
                            
            return