# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

class ButterflySpreadAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 4, 1)
        self.SetEndDate(2017, 5, 30)
        self.SetCash(150000)
        equity = self.AddEquity("GOOG", Resolution.Minute)
        self.underlying = equity.Symbol
        self.SetBenchmark(equity.Symbol)
        self.itm_call = str()
        self.otm_call = str()
        self.atm_call = str()

    def OnData(self,slice):

        ''' OptionChainProvider gets the option chain provider,
            used to get the list of option contracts for an underlying symbol.
            Then you can manually filter the contract list returned by GetOptionContractList.
            The manual filtering will be limited to the information
            included in the Symbol (strike, expiration, type, style) and/or prices from a History call '''
        if self.Securities.ContainsKey(self.itm_call) and self.Securities.ContainsKey(self.otm_call) and self.Securities.ContainsKey(self.atm_call) and not self.Portfolio.Invested:
            self.Buy(self.itm_call, 1)
            self.Buy(self.otm_call, 1) 
            self.Sell(self.atm_call, 2)
        
        if self.Portfolio[self.underlying].Invested:
            self.Liquidate(self.underlying)

        if not self.Portfolio.Invested:

            filtered_contracts = self.InitialFilter(-9, 9, 30, 60)
            # sorted the optionchain by expiration date and choose the furthest date
            if len(filtered_contracts) < 1: return
            expiry = sorted(filtered_contracts,key = lambda x: x.ID.Date, reverse=True)[0].ID.Date
            # filter the call options from the contracts expires on that date
            call = [i for i in filtered_contracts if i.ID.Date == expiry and i.ID.OptionRight == 0]
            # sorted the contracts according to their strike prices 
            call_contracts = sorted(call,key = lambda x: x.ID.StrikePrice)    
            # choose OTM call 
            self.otm_call = call_contracts[-1]
            # choose ITM call 
            self.itm_call = call_contracts[0]
            # choose ATM call
            undelyingPrice = self.Securities[self.underlying].Price
            self.atm_call = sorted(call_contracts,key = lambda x: abs(undelyingPrice - x.ID.StrikePrice))[0]
            self.trade_contracts = [self.itm_call, self.otm_call, self.atm_call]
            for i in self.trade_contracts:        
                self.AddOptionContract(i, Resolution.Minute)


    def InitialFilter(self, min_strike_rank, max_strike_rank, min_expiry, max_expiry):
        
        ''' This method is an initial filter of option contracts 
            according to the range of strike price and the expiration date '''
            
        contracts = self.OptionChainProvider.GetOptionContractList(self.underlying, self.Time.date())
        if len(contracts) == 0 : return []
        # fitler the contracts based on the expiry range
        contract_list = [i for i in contracts if min_expiry < (i.ID.Date.date() - self.Time.date()).days < max_expiry]
        # find the strike price of ATM option
        atm_strike = sorted(contract_list,
                            key = lambda x: abs(x.ID.StrikePrice - self.Securities[self.underlying].Price))[0].ID.StrikePrice
        strike_list = sorted(set([i.ID.StrikePrice for i in contract_list]))
        # find the index of ATM strike in the sorted strike list
        atm_strike_rank = strike_list.index(atm_strike)
        try: 
            strikes = strike_list[(atm_strike_rank + min_strike_rank):(atm_strike_rank + max_strike_rank)]
        except:
            strikes = strike_list
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice in strikes]
        return filtered_contracts 

    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))