from AlgorithmImports import *

class PutCalendarSpreadStrategy(QCAlgorithm): 
    def initialize(self):
        self.set_start_date(2017, 2, 1)
        self.set_end_date(2017, 2, 20)
        self.set_cash(500000)

        option = self.add_option("GOOG", Resolution.MINUTE)
        self.symbol = option.symbol
        option.set_filter(self.universe_func)

    def universe_func(self, universe):
        return universe.include_weeklys().put_calendar_spread(0, 30, 60)

    def on_data(self, data):
        # avoid extra orders
        if self.portfolio.invested: return

        # Get the OptionChain of the self.symbol
        chain = data.option_chains.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.put_calendar_spread(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 on_end_of_algorithm(self):
        for symbol, sec in self.securities.items():
            self.log(f"{symbol} :: {sec.price}")