from AlgorithmImports import *

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

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

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

    def OnData(self, data):
        # avoid extra orders
        if self.Portfolio.Invested: return

        # Get the OptionChain of the self.symbol
        chain = data.OptionChains.get(self.symbol, None)
        if not chain: return

        # get at-the-money strike
        atm_strike = sorted(chain, key=lambda x: abs(x.Strike - chain.Underlying.Price))[0].Strike

        # filter the put options from the contracts which is ATM in the option chain.
        puts = [i for i in chain if i.Strike == atm_strike and i.Right == OptionRight.Put]
        if len(puts) == 0: return

        # sorted the optionchain by expiration date
        expiries = sorted([x.Expiry for x in puts], key = lambda x: x)
        
        # select the farest expiry as far-leg expiry, and the nearest expiry as near-leg expiry
        near_expiry = expiries[0]
        far_expiry = expiries[-1]

        option_strategy = OptionStrategies.PutCalendarSpread(self.symbol, atm_strike, near_expiry, far_expiry)
        # We open a position with 1 unit of the option strategy
        self.Buy(option_strategy, 1)
        # self.Sell(option_strategy, 1) if short put calendar spread
        
    def OnEndOfAlgorithm(self):
        for symbol, sec in self.Securities.items():
            self.Log(f"{symbol} :: {sec.Price}")