from datetime import timedelta
from clr import AddReference
# import numpy as np
# from QuantConnect.Securities.Option import OptionPriceModels

class CondorAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 2, 1)
        self.SetEndDate(2017, 3, 1)
        self.SetCash(500000)
        equity = self.AddEquity("GOOG", Resolution.Minute)
        equity_cost = self.AddEquity("COST", Resolution.Minute)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.underlyingsymbol = equity.Symbol
        self.underlyingsymbol_cost = equity_cost.Symbol
        self.SetBenchmark(equity.Symbol)
        self.SetBenchmark(equity_cost.Symbol)
        
        # Jing's code:
        # self.Securities["GOOG"].MarginModel =  PatternDayTradingMarginModel()
        # self.Securities["GOOG"].SetLeverage(2)
        # option = self.AddOption("GOOG")
        # option.SetFilter(-10, +10, timedelta(0), timedelta(30))
        # option.PriceModel = OptionPriceModels.CrankNicolsonFD()
        # self.SetWarmUp(TimeSpan.FromDays(7))

    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 '''
        # slice.Contains(symbol)
        if self.Portfolio[self.underlyingsymbol].Quantity != 0 and self.Portfolio[self.underlyingsymbol_cost].Quantity != 0:
            self.Liquidate()
            
      
        if not self.Portfolio.Invested:
            contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
            self.TradeOptions(contracts)

                        
    def TradeOptions(self,contracts):
        if len(contracts) == 0 : return
        filtered_contracts = self.InitialFilter(self.underlyingsymbol, contracts, -50, 50, 0, 7)
        # sorted the optionchain by expiration date and choose the furthest date
        expiry = sorted(filtered_contracts,key = lambda x: x.ID.Date)[-1].ID.Date

        # filter the call and put options from the contracts
        call = [i for i in filtered_contracts if i.ID.OptionRight == 0 and i.ID.Date == expiry]
        # put = [i for i in filtered_contracts if i.ID.OptionRight == 1 and i.ID.Date == expiry]
        # sorted the contracts according to their strike prices 
        call_contracts = sorted(call,key = lambda x: x.ID.StrikePrice)    
        # put_contracts = sorted(put,key = lambda x: x.ID.StrikePrice)
        
        
        
        # my long condor components
        itm_call_lower = call_contracts[-15]
        itm_call_upper = call_contracts[-10]
        otm_call_lower = call_contracts[-5]
        otm_call_upper = call_contracts[-1]
        
        if itm_call_lower == 0:
            return
        if itm_call_upper == 0:
            return
        if otm_call_lower == 0:
            return
        if otm_call_upper == 0:
            return
        
        print(itm_call_lower)
        
        
        # QC's iron condor components
        # otm_put_lower = put_contracts[0]
        # otm_put = put_contracts[10]
        # otm_call = call_contracts[-10]
        # otm_call_higher = call_contracts[-1]
        
        self.trade_contracts = [itm_call_lower,itm_call_upper,otm_call_lower,otm_call_upper]
        for contract in self.trade_contracts:
            self.AddOptionContract(contract, Resolution.Minute) 
      
        self.Buy(itm_call_lower, 1) # Buy 1 ITM Call
        self.Sell(itm_call_upper, 1) # Sell 1 ITM Call 
        self.Sell(otm_call_lower, 1) # Sell 1 OTM Call 
        self.Buy(otm_call_upper, 1) # Buy 1 OTM Call
        
    
        
            
    def InitialFilter(self, underlyingsymbol, symbol_list, min_strike_rank, max_strike_rank, min_expiry, max_expiry):
        
        ''' This method is an initial filter of option contracts
            based on the range of strike price and the expiration date '''
            
        if len(symbol_list) == 0 : return
        # fitler the contracts based on the expiry range
        contract_list = [i for i in symbol_list 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[underlyingsymbol].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: 
            min_strike = strike_list[atm_strike_rank + min_strike_rank + 1]
            max_strike = strike_list[atm_strike_rank + max_strike_rank - 1]

        except:
            min_strike = strike_list[0]
            max_strike = strike_list[-1]
           
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice >= min_strike and i.ID.StrikePrice <= max_strike]

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