from AlgorithmImports import *

class PutCalendarSpreadStrategy(QCAlgorithm): 
    def Initialize(self):
        self.SetStartDate(2017, 2, 1)
        self.SetEndDate(2017, 2, 1)
        self.SetCash(500000)

        option = self.AddOption("GOOG", Resolution.Minute)
        self.symbol = option.Symbol
        option.SetFilter(self.UniverseFunc)

        self.step = 1

    def UniverseFunc(self, universe):
        return universe.Strikes(-1, 1).Expiration(timedelta(0), timedelta(62))

    def OnData(self, data):
        # we will only execute the 2 steps: 
        # 1. sell put in the front month
        # 2. roll put out to back month
        if self.step > 2: return

        chain = data.OptionChains.get(self.symbol, None)
        if not chain: return

        if self.step == 1:
            self.atm_strike = sorted(chain, key=lambda x: abs(x.Strike - chain.Underlying.Price))[0].Strike

        puts = [i for i in chain if i.Strike == self.atm_strike and i.Right == OptionRight.Put]
        if len(puts) == 0: return

        # step 1: sell 1 atm put in the front month
        if self.step == 1:
            # front month contract
            put = sorted(puts, key = lambda x: x.Expiry)[0]
            self.Sell(put.Symbol, 1)
            self.step += 1

        # step 2: attempt to roll the put out to back month - error occurs: 
        # Error executing margin models: Unable to create group for orders: [2,3]
        elif self.step == 2:
            expiries = sorted([x.Expiry for x in puts], key = lambda x: x)
            
            near_expiry = expiries[0] # front month
            far_expiry = expiries[-1] # back month
            option_strategy = OptionStrategies.PutCalendarSpread(self.symbol, self.atm_strike, near_expiry, far_expiry)
            self.Sell(option_strategy, 1)
            self.step += 1

    def OnEndOfAlgorithm(self):
        for x in self.Portfolio:
            if (x.Value.Invested):
                self.Log(f"{x.Key} :: {x.Value.Quantity}")

    def OnOrderEvent(self, orderEvent: OrderEvent) -> None:
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if orderEvent.Status == OrderStatus.Filled:
            self.Debug(f"{self.Time}: {order.Type}: {orderEvent}")