# region imports
from AlgorithmImports import *
# endregion

symbol_dict = {  # will be used to store the configs of multiple symbols
                'SPY':  # equity symbol
                    {
                        'initialAmount': 10000,  # adjust this for trade size
                        'contractDelta': 0.5,  # sets the delta for contract to select
                        'SMA': 25,  # SMA for the indicator
                        'waitPeriod': 0,  # additional days to wait after exit
                        'trailingStop': 0.03  # trailing stop percentage
                    }
                }

class SMAOptions(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2022, 1, 1)  # Set Start Date 
        self.SetEndDate(2022, 6, 30)  # sets end date
        self.SetCash(1000000)  # Set starting account size

        # settings for option contracts expirations
        self.expirationThreshold = timedelta(10)  # this is the days threshold to perform a time roll
        self.minContractExpiration = timedelta(30)  # this is the minimum number of days out to look for a contract
        self.maxContractExpiration = timedelta(60)  # this is the maximum number of days out to look for a contract

        # sets brokerage model
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # setting empty dictionaries for multiple symbols
        self.tradeAmount = {}
        self.isRollOut = {}
        self.rollValue = {}
        self.equity = {}
        self.equityDailyPrice = {}
        self.equitySingleDayPrice = {}
        self.entryAmount = {}
        self.entryStrike = {}
        self.manualSMA1 = {}
        self.manualSMA2 = {}
        self.autoSMA = {}
        self.waitPeriod = {}
        self.nextEntryTime = {}
        self.nextExitTime = {}
        self.trailingStop = {}
        self.exitStop = {}
        self.call = {}

        # initialize a warm up period of 0
        warmUpPeriod = 0
        
        self.exitPeriod = timedelta(0.5) 

        # loop to initialize each symbol
        for symbol in symbol_dict.keys():
            
            # adds config parameters to log
            symbolConfigMessage = "[Symbol] {0}, [initialAmount] {1}, [contractDelta] {2}, [SMA] {3}, [trailingStop] {4}".format(
                                        symbol,
                                        symbol_dict[symbol]['initialAmount'],
                                        symbol_dict[symbol]['contractDelta'],
                                        symbol_dict[symbol]['SMA'],
                                        symbol_dict[symbol]['waitPeriod'],
                                        symbol_dict[symbol]['trailingStop']
                                    )
            self.Log(symbolConfigMessage)

            self.tradeAmount[symbol] = 0
            
            self.isRollOut[symbol] = False
            self.rollValue[symbol] = 0

            # sets underlying asset
            self.equity[symbol] = self.AddEquity(symbol, Resolution.Minute)  
            self.equity[symbol].SetDataNormalizationMode(DataNormalizationMode.Raw)
            
            self.equityDailyPrice[symbol] = None
            self.entryAmount[symbol] = 0  # used to track initial entry over rolls for pnl
            self.entryStrike[symbol] = 0. # used to track initial strike price

            option = self.AddOption(self.equity[symbol].Symbol, Resolution.Minute) 
            option.SetFilter(-2, 2, self.minContractExpiration, self.maxContractExpiration)
            # set pricing model to calculate Greeks
            option.PriceModel = OptionPriceModels.BjerksundStensland()

            # initialize indicators
            self.manualSMA1[symbol] = SimpleMovingAverage(self.equity[symbol].Symbol, symbol_dict[symbol]['SMA'])
            self.manualSMA2[symbol] = SimpleMovingAverage(self.equity[symbol].Symbol, symbol_dict[symbol]['SMA'])
            self.autoSMA[symbol] = self.SMA(self.equity[symbol].Symbol, symbol_dict[symbol]['SMA'], Resolution.Daily)

            self.waitPeriod[symbol] = timedelta(0.5) + timedelta(symbol_dict[symbol]['waitPeriod'])
            self.nextEntryTime[symbol] = self.Time
            self.nextExitTime[symbol] = self.Time

            # trailing stop initialization
            self.trailingStop[symbol] = 0
            self.exitStop[symbol] = 0

            # calculates warm up period to be maximum SMA out of the equities
            warmUpPeriod = max(warmUpPeriod, symbol_dict[symbol]['SMA'])

            # needed for plotting each symbol
            stockPlot = Chart(f'{symbol} Plot')
            stockPlot.AddSeries(Series('Price', SeriesType.Candle, '$'))
            self.AddChart(stockPlot)

        # Generating consolidators
        consDaily = TradeBarConsolidator(Resolution.Daily)
        consDaily.DataConsolidated += self.OnDailyData
        self.SubscriptionManager.AddConsolidator(self.equity['SPY'].Symbol, consDaily)

        consSingleDay = TradeBarConsolidator(timedelta(1))
        consSingleDay.DataConsolidated += self.OnSingleDayData 
        self.SubscriptionManager.AddConsolidator(self.equity['SPY'].Symbol, consSingleDay)

        self.exchange = self.Securities[self.equity[symbol].Symbol].Exchange  # using to check when exchange is open

        # sets the warm up period
        self.SetWarmUp(timedelta(warmUpPeriod))

        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(9, 32), self.PlotLogDaily)


    def OnData(self, data: Slice):
        # checks if smas are ready or if no daily price
        if self.IsWarmingUp:
            return

        for symbol in symbol_dict.keys():
            #if not self.daily_bb[symbol].IsReady:
            #    continue
            if self.equityDailyPrice[symbol] is None:
                continue

            # runs only when exchange hours are open
            if self.exchange.ExchangeOpen:

                # checks for open option positions
                option_invested = [x for x, holding in self.Portfolio.items() if holding.Invested and holding.Type == SecurityType.Option and x.HasUnderlying and x.Underlying.Equals(self.equity[symbol].Symbol)]
                
                if option_invested:
                    # calculates trailing and exit stop:
                    self.trailingStop[symbol] = self.equityDailyPrice[symbol] * (1-symbol_dict[symbol]['trailingStop'])

                    self.exitStop[symbol] = max(self.exitStop[symbol], self.trailingStop[symbol])

                    # Exits if price is below the stop
                    if self.exitStop[symbol] > self.equity[symbol].Price:
                        if self.nextExitTime[symbol] <= self.Time:  # exit needs to wait a day after entry
                            pnl = self.Securities[option_invested[0]].Holdings.HoldingsValue - self.entryAmount[symbol]

                            orderMessage = "[Time] {0}; [Underlying] {1}; [Strike Price] {2}; [Contract Delta] {3}; [Contract Expiration] {4}; [Enter Trigger] {5:.2f}; [Enter Amount] {6:.2f}; [Portfolio Value] {7:.2f}; [Trailing Stop] {8:.2f}; [Exit Stop] {9:.2f}; [PNL] {10:.2f}".format(
                                    self.Time,
                                    self.equity[symbol].Price,
                                    self.call[symbol].Strike,
                                    self.call[symbol].Greeks.Delta,
                                    self.call[symbol].Expiry,
                                    self.autoSMA[symbol].Current.Value,
                                    self.entryAmount[symbol],
                                    self.Portfolio.TotalPortfolioValue,
                                    self.trailingStop[symbol],
                                    self.exitStop[symbol],
                                    pnl
                                )
                            self.Liquidate(option_invested[0], orderMessage)
                            self.Log(orderMessage)
                            self.nextEntryTime[symbol] = self.Time + self.waitPeriod[symbol]  # Waits at least until next day to re-enter

                            # resets all flags/counters
                            self.trailingStop[symbol] = 0
                            self.exitStop[symbol] = 0
                            self.entryAmount[symbol] = 0
                            self.entryStrike[symbol] = 0
                    
                    # If too close to expiration, liquidate and perform a time roll
                    elif self.Time + self.expirationThreshold > option_invested[0].ID.Date:
                        self.rollValue[symbol] = self.Securities[option_invested[0]].Holdings.HoldingsValue 
                        timeRollMessage = "[Time] {0}; [Contract Expiration] {1}; [Time Roll Amount] {2}".format(
                                            self.Time,
                                            self.call[symbol].Expiry,
                                            self.rollValue[symbol]
                                            )
                        self.Liquidate(option_invested[0], timeRollMessage)
                        self.Log(timeRollMessage)
                        self.isRollOut[symbol] = True

                        # determines the option chain and buys the option
                        for i in data.OptionChains:
                            chains = i.Value
                            self.BuyCall(chains, symbol)
                    
                # Attempts entry if no open options positions
                else:
                    # entry if above upper band
                    if self.autoSMA[symbol].Current.Value <= self.equity[symbol].Price:
                        if self.nextEntryTime[symbol] <= self.Time:
                            self.nextExitTime[symbol] = self.Time + self.exitPeriod
                            self.entryStrike[symbol] = self.autoSMA[symbol].Current.Value
                            for i in data.OptionChains:
                                chains = i.Value
                                self.BuyCall(chains, symbol)
                        
    
    def BuyCall(self, chains, symbol):
        # sorts the options chains and picks the furthest away between 30-60 days
        expiry = sorted(chains, key = lambda x: x.Expiry, reverse=True)[0].Expiry
        # filter out calls
        calls = [i for i in chains if i.Expiry == expiry and i.Right == OptionRight.Call]

        # finds contract closest to the requested delta with and without roll
        if self.isRollOut[symbol] is False:
            call_contracts = sorted(calls, key = lambda x: abs(x.Greeks.Delta - symbol_dict[symbol]['contractDelta']))
        elif self.isRollOut[symbol]:
            call_contracts = sorted(calls, key = lambda x: abs(x.Strike - self.entryStrike[symbol]))
        else:
            self.Debug("Roll flag error")
            return

        # saves first element to the self.call variable
        if len(call_contracts) == 0:
            return
        self.call[symbol] = call_contracts[0]
        
        try:
            # calculates trade amount
            if self.isRollOut[symbol] is False:
                self.tradeAmount[symbol] = symbol_dict[symbol]['initialAmount']
            else:
                self.tradeAmount[symbol] = self.rollValue[symbol]

            # determines quantity based on trade amount
            quantity = int((self.tradeAmount[symbol]) / self.call[symbol].AskPrice / 100)

            # updates trade amount based on quantity calculated
            self.tradeAmount[symbol] = quantity * self.call[symbol].AskPrice * 100

            if self.isRollOut[symbol] is False:
                self.entryAmount[symbol] = self.tradeAmount[symbol]

            buyLogMessage= "[Time] {0}; [Underlying] {1}; [Strike Price] {2}; [Contract Delta] {3}; [Contract Expiration] {4}; [Enter Trigger] {5:.2f}; [Enter Amount] {6:.2f}; [Portfolio Value] {7:.2f}".format(
                            self.Time,
                            self.equity[symbol].Price,
                            self.call[symbol].Strike,
                            self.call[symbol].Greeks.Delta,
                            self.call[symbol].Expiry,
                            self.autoSMA[symbol].Current.Value,
                            self.tradeAmount[symbol],
                            self.Portfolio.TotalPortfolioValue
                            )
            #self.Buy(self.call.Symbol, quantity)
            self.MarketOrder(self.call[symbol].Symbol, quantity, False, buyLogMessage)
            self.Log(buyLogMessage)
        except:
            self.Debug("error computing quantity") # this is done since I had run into some errors computing the quantity



    def OnOrderEvent(self, orderEvent):
        # log order events
        self.Log("[OnOrderEvent] " + str(orderEvent))
        pass


    def OnDailyData(self, sender, bar):
        # updates indicator calculations and daily price on consolidated daily bar
        self.manualSMA1['SPY'].Update(bar.EndTime, bar.Close)
        self.equityDailyPrice['SPY'] = bar.Close
    

    def OnSingleDayData(self, sender, bar):
        self.manualSMA2['SPY'].Update(bar.EndTime, bar.Close)
        self.equitySingleDayPrice['SPY'] = bar.Close


    def PlotLogDaily(self):
        for symbol in symbol_dict.keys():
            try:
                dailyMessage = "[Time] {0}, [Symbol] {1}, [Manual SMA Daily] ${2:.2f}, [Daily Close] ${3:.2f}, [Manual SMA Single Day] ${4:.2f}, [Single Day Close] ${5:.2f}, [Auto SMA] ${6:.2f}".format(
                                self.Time,
                                symbol,
                                self.manualSMA1[symbol].Current.Value,
                                self.equityDailyPrice[symbol],
                                self.manualSMA2[symbol].Current.Value,
                                self.equitySingleDayPrice[symbol],
                                self.autoSMA[symbol].Current.Value
                                )
                self.Log(dailyMessage)    
                # plots the daily price and the entry line
                self.Plot(f'{symbol} Plot', 'Price', self.equityDailyPrice[symbol])
                self.Plot(f'{symbol} Plot', 'SMA', self.autoSMA[symbol].Current.Value)
            except:
                self.Debug("symbol not found")