from scipy.stats import norm
from datetime import timedelta

class IronCondorAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 2, 1)
        self.SetEndDate(2018, 2, 1)
        self.SetCash(100000)
        self.Debug('Starting Cash: ' + str(self.Portfolio.Cash))
        # Add equities
        self.AddUniverse(self.Universe.DollarVolume.Top(50))
        self.equity = [x for x in self.ActiveSecurities]
        self.equity.append(self.AddEquity("SPY", Resolution.Minute))
        self.equity.append(self.AddEquity("EEM", Resolution.Minute))
        self.equity.append(self.AddEquity("GOOG", Resolution.Minute))
        self.equity.append(self.AddEquity("VCR", Resolution.Minute))
        self.equity.append(self.AddEquity("SLY", Resolution.Minute))
        self.equity.append(self.AddEquity("XLP", Resolution.Minute))
        self.equity.append(self.AddEquity("ARKK", Resolution.Minute))
        self.equity.append(self.AddEquity("XLY", Resolution.Minute))
        self.equity.append(self.AddEquity("AAPL", Resolution.Minute))
        self.stock_list = ["GOOG", "SPY", "VCR", "SLY", "EEM", "XLP", "ARKK", "XLY", "AAPL"]
        for i in self.equity:
            i.SetLeverage(1)
        self.options = []
        # Create an array of options for each symbol
        for x in self.equity:
            self.options.append(self.AddOption(x.Symbol, Resolution.Minute))
        # Create symbol array
        # self.symbols = [x for x.Symbol in self.options]
        # Specify universe function
        for x in self.options:
            x.SetFilter(self.UniverseFunc)
        # use the underlying equity GOOG as the benchmark
        self.SetBenchmark("SPY")
        self.counter = 0
        self.liabilities = 0
        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)

        
    # What I think is happening is it's sending a ton of data simultaneously
    def OnData(self, slice):
        self.most_recent_slice = slice

    def TradeOptions(self):
        # If there is undelying assets in portfolio at expiration, liquidate the stocks in order to roll into new contracts
        # (Fixed) PROBLEM: is datetime retrieving now-now or backtest-now
        self.Debug(str(x for x in self.Portfolio))
        #This is a problem; it needs to be taylored to each option
        if self.Portfolio.Invested:
            self.AdjustOptions()
        if not self.Portfolio.Invested and self.Time.hour != 0 and self.Time.minute != 0: 
            # List of optionchains for each security
            self.AddToPortfolio()
    
    def CloseCondor(self, x):
        #CCC = Close Condor Called
        self.Debug("CCC")
        for symbol in x:
            self.Liquidate(symbol)
    
    def AdjustOptions(self):
        if self.Portfolio.Invested:
            # self.Debug("Portfolio.Invested Called " + str(self.Time))
            for x in self.condor_list:
                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: #put
                        if c.UnderlyingLastPrice - c.Strike < 0:
                            otm = False
                    else:
                        if c.UnderlyingLastPrice - c.Strike > 0:
                            otm = False
                # If Condor is in the head region, close
                if days_to_expiry <= 25 and otm:
                    self.CloseCondor(x)
    
            #Clean out actual stocks
            for ticker in self.stock_list:
                self.SetHoldings(ticker, 0)
    
            self.AddToPortfolio()
        
    def AddToPortfolio(self):
        self.condor_list = []
        for i in self.most_recent_slice.OptionChains:
            # self.Debug("AddToPortfolio Called")
            # self.Debug("Iteration Number " + str(self.counter))
            chain = i.Value
            contract_list = [x for x in chain]
            # if there is no optionchain or no contracts in this optionchain, pass the instance
            if (self.most_recent_slice.OptionChains.Count == 0) or (len(contract_list) == 0): 
                continue   

            # sorted the optionchain by expiration date and choose the furthest date
            expiry = sorted(chain,key = lambda x: x.Expiry)[-1].Expiry
            
            # filter the call and put options from the contracts
            
            put = [i for i in chain if i.Expiry == expiry and i.Right == 1]
            self.Debug("Puts: " + str(x.Strike for x in put))
            if len(put) > 2:
                put_contracts = sorted(put,key = lambda x: x.Strike)
            else:
                continue
            highest_put_strike = put_contracts[len(put_contracts) - 1].Strike
            
            call = [i for i in chain if i.Expiry == expiry and i.Right == 0 and i.Strike > put[0].UnderlyingLastPrice]
            if len(call) > 2:
                call_contracts = sorted(call,key = lambda x: x.Strike) 
            else:
                continue
            
            # sorted the contracts according to their strike prices 
            self.Debug(put_contracts[len(put_contracts) - 1].UnderlyingSymbol.Value + ' ' + str(len(put_contracts)) + ' ' + str(len(call_contracts)))
            # self.Debug('price= ' + str(put_contracts[0].UnderlyingLastPrice) + 'PUTS: ' + ' '.join(str([x.Strike for x in put_contracts])))
            put_contracts = [i for i in put if i.Strike < i.UnderlyingLastPrice]
            # self.Debug('CALLS: ' + ' '.join([x.Strike for x in call_contracts]))

            put_contracts = sorted(put,key = lambda x: x.Strike)    
            if len(call_contracts) == 0 or len(put_contracts) == 0 : continue

            otm_put_lower = put_contracts[0]
            otm_put = put_contracts[len(put_contracts) - 1]
            try:
                otm_call = call_contracts[-10]
            except IndexError:
                otm_call = call_contracts[-1 * len(call_contracts)]
                self.Debug("otm_call Index Error")
            try:
                otm_call_higher = call_contracts[-1]
            except IndexError:
                otm_call_higher = call_contracts[1]
                self.Debug("otm_call_higher Index Error")

            self.trade_contracts = [otm_call.Symbol,otm_call_higher.Symbol,otm_put.Symbol,otm_put_lower.Symbol]
            self.condor_list.append([otm_call, otm_call_higher, otm_put, otm_put_lower])
            
            self.Debug('Lower Put ' + str(otm_put_lower.Strike) + 'Higher Put ' + str(otm_put.Strike) + 'Lower Call: ' + str(otm_call.Strike) + 'Higher Call' + str(otm_call_higher.Strike))
            expiry = otm_call_higher.Expiry
            # if there is no securities in portfolio, trade the options 
            totalPrice = sum([x.AskPrice for x in [otm_call_higher, otm_put_lower]]) - otm_put.BidPrice - otm_call.BidPrice + otm_call.UnderlyingLastPrice
            margin = self.Portfolio.GetMarginRemaining(otm_put_lower.UnderlyingSymbol, OrderDirection.Buy)
            if margin > totalPrice * 4:
                self.Debug("Passed Enough Cash Test. Cash: " + str(self.Portfolio.Cash))
                self.Debug("Option Total Pricetag: " + str(totalPrice))
                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)
                # Fix, this doesn't include premium
            self.counter += 1
            self.liabilities += 100 * otm_call.UnderlyingLastPrice - sum([x.AskPrice for x in [otm_call_higher, otm_put_lower]]) - otm_put.BidPrice - otm_call.BidPrice

    def ExpectedValue(x):
        # What will q be? 100? (that is assumed here)
        mu = x[0].UnderlyingLastPrice
        sigma = x[0].ImpliedVolatility
        C = x[1].BidPrice + x[2].BidPrice - x[0].BidPrice - x[3].BidPrice
        y = []
        y[0] = norm.cdf(x[0].Strike, mu, sigma) * -100 * (x[1].Strike - x[0].Strike)
        Ptwo = norm.cdf(x[1].Strike, mu, sigma) - norm.cdf(x[0].Strike, mu, sigma)
        y[1] = Ptwo * -100 
        y[2] = 2
        y[3] = 3
        y[4] = 4
        
    def OnOrderEvent(self, orderEvent):
        self.Debug(str(orderEvent))

    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys().Strikes(-15, 15).Expiration(timedelta(35), timedelta(50))