from AlgorithmImports import *

class PutButterflyStrategy(QCAlgorithm): 
    def initialize(self):
        self.set_start_date(2017, 2, 1)
        self.set_end_date(2017, 3, 6)
        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_butterfly(30, 5)

    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

        # sorted the optionchain by expiration date and choose the furthest date
        expiry = sorted(chain, key = lambda x: x.expiry, reverse=True)[0].expiry
        
        # filter the put options from the contracts which expire on the furthest expiration date in the option chain.
        puts = [i for i in chain if i.expiry == expiry and i.right == OptionRight.PUT]
        if len(puts) == 0: return

        # sort the put options with the same expiration date according to their strike price.
        put_strikes = sorted([x.strike for x in puts])

        # get at-the-money strike
        atm_strike = sorted(puts, key=lambda x: abs(x.strike - chain.underlying.price))[0].strike

        # Get the distance between lowest strike price and ATM strike, and highest strike price and ATM strike. 
        # Get the lower value as the spread distance as equidistance is needed for both side.
        spread = min(abs(put_strikes[0] - atm_strike), abs(put_strikes[-1] - atm_strike))

        # select the strike prices for forming the option legs
        itm_strike = atm_strike + spread
        otm_strike = atm_strike - spread

        option_strategy = OptionStrategies.short_butterfly_put(self.symbol, itm_strike, atm_strike, otm_strike, 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 butterfly