from scipy.stats import norm
from datetime import timedelta

class IronCondorAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 2, 1)
        self.SetEndDate(2017, 3, 1)
        self.SetCash(100000)
        
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        
        # Add equities
        self.AddUniverse(self.Universe.DollarVolume.Top(5))
        [self.AddEquity(symbol, Resolution.Minute) for symbol in ["GOOG", "SPY", "VCR", "SLY", "EEM", "XLP", "ARKK", "XLY", "AAPL"]]
        self.activeSymbols = []
        
        # use the underlying equity GOOG as the benchmark
        self.SetBenchmark("SPY")
        self.SetWarmUp(1)
        # Set trading intervals
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("SPY", 1), self.TradeOptions)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(12, 0), self.TradeOptions)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose("SPY", 10), self.TradeOptions)
    
    def CustomSecurityInitializer(self, security):
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)
        security.SetMarketPrice(self.GetLastKnownPrice(security))
        
    def OnSecuritiesChanged(self, changes):
        for change in changes.AddedSecurities:
            if change.Type == SecurityType.Equity:
                change.SetLeverage(1)
                self.activeSymbols.append(change.Symbol)
                
        for change in changes.RemovedSecurities:
            if change.Symbol in self.activeSymbols:
                self.activeSymbols.remove(change.Symbol)
                
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if order.Type == OrderType.OptionExercise:
            self.Liquidate(orderEvent.Symbol.Underlying)

    def TradeOptions(self):
        if self.Portfolio.Invested:
            self.AdjustOptions()
        if not self.Portfolio.Invested and self.Time.hour != 0 and self.Time.minute != 0: 
            self.AddToPortfolio()
    
    def CloseCondor(self, x):
        #CCC = Close Condor Called
        self.Debug("CCC")
        for symbol in x:
            self.Liquidate(symbol.Symbol)
    
    def AdjustOptions(self):
        if self.Portfolio.Invested:
            for symbol, x in self.condor_list.items():
                days_to_expiry = abs(x[0].Expiry - self.Time).days
                # if this condor expires in 25+ days, leave it condor
                if days_to_expiry > 25:
                    continue
                elif days_to_expiry < 3.75:
                    self.CloseCondor(x)
                    
                # OTM Check
                otm = True
                for c in x:
                    if c.Right == 1 and self.Securities[symbol].Price - c.StrikePrice < 0:
                        otm = False
                    elif c.Right == 0 and self.Securities[symbol].Price - c.StrikePrice > 0:
                        otm = False
                # If Condor is in the head region, close
                if days_to_expiry <= 25 and otm:
                    self.CloseCondor(x)
        
    def AddToPortfolio(self):
        self.condor_list = {}
        
        for symbol in self.activeSymbols:
            contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time)
            # option universe filter
            contracts = self.InitialFilter(symbol, contracts, -15, 15, 35, 50)
            if not contracts: continue
            
            # sorted the optionchain by expiration date and choose the furthest date
            expiry = sorted(contracts, key = lambda x: x.ID.Date, reverse=True)[0].ID.Date
            
            # get puts
            puts = [contract for contract in contracts if contract.ID.Date == expiry \
                                                        and contract.ID.OptionRight == OptionRight.Put \
                                                        and contract.ID.StrikePrice < self.Securities[symbol].Price]
            if len(puts) <= 2: continue
            put_contracts = sorted(puts, key = lambda x: x.ID.StrikePrice)
            
            # get calls
            calls = [contract for contract in contracts if contract.ID.Date == expiry \
                                                        and contract.ID.OptionRight == OptionRight.Call \
                                                        and contract.ID.StrikePrice > self.Securities[symbol].Price]
            if len(calls) <= 2: continue
            call_contracts = sorted(calls, key = lambda x: x.ID.StrikePrice) 
            
            # get and subscribe to the option contracts
            otm_put_lower = self.AddOptionContract(put_contracts[0], Resolution.Minute)
            otm_put = self.AddOptionContract(put_contracts[-1], Resolution.Minute)
            otm_call = self.AddOptionContract(call_contracts[0], Resolution.Minute)
            otm_call_higher = self.AddOptionContract(call_contracts[-1], Resolution.Minute)
            self.condor_list[symbol] = [otm_call, otm_call_higher, otm_put, otm_put_lower]
            
            # get the margin requirement
            totalPrice = (otm_call_higher.AskPrice + otm_put_lower.AskPrice + otm_put.BidPrice + otm_call.BidPrice) * 100
            margin = self.Portfolio.GetMarginRemaining(symbol, OrderDirection.Buy)
            if margin > totalPrice * 4:
                self.Buy(otm_put_lower.Symbol ,1)
                self.Sell(otm_put.Symbol ,1)
                self.Sell(otm_call.Symbol ,1)
                self.Buy(otm_call_higher.Symbol ,1)
                
    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 
            https://www.quantconnect.com/tutorials/applied-options/iron-condor'''
        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