# region imports
from AlgorithmImports import *
# endregion

# Import datetime
import datetime

class DancingBrownSheep(QCAlgorithm):

    def Initialize(self):

        # # # User input section # # #

        # Set Start Date
        self.SetStartDate(2021, 3, 1)

        # Set End Date
        self.SetEndDate(2021, 4, 1)

        # Set Strategy Cash
        self.SetCash(100000)

        # Timeframe (in minutes)
        time_frame = 5

        # List of stocks
        list_of_stocks = ["SPY"]

        # EMA 1
        EMA_one = 9

        # EMA 2
        EMA_two = 13

        # EMA 3
        EMA_three = 21

        # EMA 4
        EMA_four = 200

        # RSI
        RSI_period = 14

        # Warmup days
        warmup_days = 10

        # # # End # # # 

        # Symbols dictionary
        self.symbols_list = []

        # EMA 10 dictionary
        self.EMA_ten_dict = {}

        # EMA 20 dictionary
        self.EMA_twenty_dict = {}

        # EMA 30 dictionary
        self.EMA_thirty_dict = {}

        # EMA 40 dictionary
        self.EMA_forty_dict = {}

        # RSI 14 dictionary
        self.RSI_dict = {}

        # Previous EMA storage
        self.previous_ema_one_ema_three_difference = {}

        # Rolling RSI dictionary
        self.rolling_rsi_dict = {}

        # Rolling bars storage
        self.rolling_bars_storage = {}

        # Call tradable list for signal 1
        self.call_signal_one_tradable_list = []

        # Call tradable list for signal 2 - 1
        self.call_signal_two_tradable_list_one = []

        # Call tradable list for signal 2 - 2
        self.call_signal_two_tradable_list_two = []

        # Put tradable list for signal 1
        self.put_signal_one_tradable_list = []

        # Put tradable list for signal 2 - 1
        self.put_signal_two_tradable_list_one = []

        # Put tradable list for signal 2 - 2
        self.put_signal_two_tradable_list_two = []

        # Counter
        self.counter = 0

        # Data storage
        self.data_storage = 0

        # Loop through list
        for ticker in list_of_stocks:

            # Initialize AAPL
            symbol = self.AddEquity(ticker, Resolution.Minute).Symbol

            # Option
            option = self.AddOption(symbol)

            # Set filter
            option.SetFilter(-2, +2, timedelta(days = 0), timedelta(days = 7))

            # Save stock symbol
            self.symbols_list.append(symbol)

            #  Create 5-minute consolidator
            five_min_cons = TradeBarConsolidator(timedelta(minutes=time_frame))

            # Assign FiveMinuteHandler as the receiver of consolidated bars for our consolidator object
            five_min_cons.DataConsolidated += self.FiveMinuteHandler
            
            # Subscribe consolidator object to be automatically updated with minute bars
            self.SubscriptionManager.AddConsolidator(symbol, five_min_cons)

            # Create EMA-10
            self.EMA_ten_dict[symbol] = ExponentialMovingAverage(EMA_one)

            # Register indicator
            self.RegisterIndicator(symbol, self.EMA_ten_dict[symbol], five_min_cons)

            # Create EMA-20
            self.EMA_twenty_dict[symbol] = ExponentialMovingAverage(EMA_two)

            # Register indicator
            self.RegisterIndicator(symbol, self.EMA_twenty_dict[symbol], five_min_cons)

            # EMA 30 dictionary
            self.EMA_thirty_dict[symbol] = ExponentialMovingAverage(EMA_three)

            # Register indicator
            self.RegisterIndicator(symbol, self.EMA_thirty_dict[symbol], five_min_cons)

            # EMA 40 dictionary
            self.EMA_forty_dict[symbol] = ExponentialMovingAverage(EMA_four)

            # Register indicator
            self.RegisterIndicator(symbol, self.EMA_forty_dict[symbol], five_min_cons)

            # RSI 14 dictionary
            self.RSI_dict[symbol] = RelativeStrengthIndex(RSI_period)

            # Register indicator
            self.RegisterIndicator(symbol, self.RSI_dict[symbol], five_min_cons)

            # Add to previous EMA difference storage
            self.previous_ema_one_ema_three_difference[symbol] = []

            # Add to rolling RSI dictionary
            self.rolling_rsi_dict[symbol] = []

            # Add to rolling bars storage
            self.rolling_bars_storage[symbol] = []

        # Warmup period
        self.SetWarmUp(timedelta(days = warmup_days))

    def OnData(self, data: Slice):

        # Save data
        self.data_storage = data.OptionChains

    def FiveMinuteHandler(self, sender, bar):

        # If not warmup
        if not self.IsWarmingUp:

            # Add 1 to counter
            self.counter += 1

            # If counter == number of symbols in algorithm
            if self.counter == len(self.symbols_list):

                # Reset counter
                self.counter = 0

            # Symbol
            symbol = bar.Symbol

            # Add bar to storage
            self.rolling_bars_storage[symbol].append(bar)

            # Add EMA difference to list
            self.previous_ema_one_ema_three_difference[symbol].append(self.EMA_ten_dict[symbol].Current.Value - self.EMA_thirty_dict[symbol].Current.Value)

            # Update RSI storage
            self.rolling_rsi_dict[symbol].append(self.RSI_dict[symbol].Current.Value)

            # If stored more than 3
            if len(self.rolling_bars_storage[bar.Symbol]) > 3:

                # Cut length
                self.rolling_bars_storage[bar.Symbol] = self.rolling_bars_storage[bar.Symbol][-3:]

                # Cut length
                self.previous_ema_one_ema_three_difference[symbol] = self.previous_ema_one_ema_three_difference[symbol][-3:]

                # Cut length
                self.rolling_rsi_dict[symbol] = self.rolling_rsi_dict[symbol][-3:]

            # If counter is 0 and stored at least 3 rolling bars
            if self.counter == 0 and len(self.rolling_rsi_dict[symbol]) >= 3:

                # Current time
                current_time = datetime.time(self.Time.hour, self.Time.minute)

                # 9:45 am time
                nine_forty_five = datetime.time(9, 45)

                # 2:45 pm time
                two_forty_five = datetime.time(14, 45)

                # If current time is between 9:45 and 2:45 pm
                if current_time >= nine_forty_five and current_time < two_forty_five:

                    # Invested
                    invested_options = []

                    # Get invested options
                    options_holdings = [x for x in self.Portfolio if x.Value.Type == SecurityType.Option and x.Value.Invested]

                    # Calls
                    stock_calls = {}

                    # Puts
                    stock_puts = {}

                    # Loop invested
                    for kvp in options_holdings:

                        # Get ID
                        option_ID = kvp.Key.ID

                        # Get value
                        option_value = kvp.Value

                        # Get underlying asset
                        underlying_asset_symbol = option_ID.Symbol

                        # Get option right
                        option_right = option_ID.OptionRight

                        # Get option symbol
                        option_symbol = option_value.Symbol

                        # If call
                        if option_right == 0:

                            # Add to call list
                            stock_calls[underlying_asset_symbol] = option_symbol
                        
                        # If put
                        elif option_right == 1:

                            # Add to put list
                            stock_puts[underlying_asset_symbol] = option_symbol

                    # Loop through stocks
                    for symbol in self.symbols_list:

                        # If value in calls list
                        if symbol.Value in stock_calls:

                            # Add to invested
                            invested_options.append(symbol)

                            # Value
                            value = symbol.Value

                            # Average price
                            average_price = self.Portfolio[stock_calls[value]].AveragePrice

                            # Current price
                            current_price = self.Securities[stock_calls[value]].AskPrice

                            # If current price 10% or lower from average price
                            if current_price <= average_price * 0.9:

                                # Close all positions
                                self.Liquidate(stock_calls[value], tag = "Stop loss - selling all holdings")

                                # Pop
                                stock_calls.pop(value)
                            
                            # Else
                            else:

                                # If close of stock less than or equal to EMA 2
                                if self.rolling_bars_storage[symbol][-1].Close <= self.EMA_twenty_dict[symbol].Current.Value:

                                    # Close all positions
                                    self.Liquidate(stock_calls[value], tag = "Call exit 2 - selling all holdings")

                                    # Pop
                                    stock_calls.pop(value)

                                # Else
                                else:

                                    # If close below EMA 10
                                    if self.rolling_bars_storage[symbol][-1].Close <= self.EMA_ten_dict[symbol].Current.Value:

                                        # Quantity
                                        quantity = self.Portfolio[stock_calls[value]].Quantity

                                        # If quantity is 1 
                                        if quantity == 1:

                                            # Sell
                                            self.Liquidate(stock_calls[value], tag = "Call exit 1 - selling all holdings")

                                            # Pop
                                            stock_calls.pop(value)

                                        # Else
                                        else:

                                            # Total cost of position
                                            cost_of_position = average_price * quantity

                                            # Value of position 
                                            value_of_position = current_price * quantity

                                            # Profit
                                            profit = value_of_position - cost_of_position

                                            # If profit greater than 0
                                            if profit > 0:

                                                # Profit / current price
                                                quantity_to_keep = profit / current_price

                                                # Quantity to sell
                                                quantity_to_sell = int(quantity - quantity_to_keep)

                                                # If quantity_to_sell greater than or equal to 1
                                                if quantity_to_sell >= 1:

                                                    # Liquidate
                                                    self.MarketOrder(stock_calls[value], -quantity_to_sell, tag = "Call exit 1 - selling breakeven")

                        # If value iin calls list
                        elif symbol.Value in stock_puts:

                            # Add to invested
                            invested_options.append(symbol)

                            # Value
                            value = symbol.Value

                            # Average price
                            average_price = self.Portfolio[stock_puts[value]].AveragePrice

                            # Current price
                            current_price = self.Securities[stock_puts[value]].AskPrice

                            # If current price 10% or lower from average price
                            if current_price <= average_price * 0.9:

                                # Close all positions
                                self.Liquidate(stock_puts[value], tag = "Stop loss - selling all holdings")

                                # Pop
                                stock_puts.pop(value)

                            # Else
                            else:

                                # If close of stock greater than or equal to EMA 2
                                if self.rolling_bars_storage[symbol][-1].Close >= self.EMA_twenty_dict[symbol].Current.Value:

                                    # Close all positions
                                    self.Liquidate(stock_puts[value], tag = "Put exit 2 - selling all holdings")

                                    # Pop
                                    stock_puts.pop(value)

                                # Else
                                else:

                                    # If close greater EMA 10
                                    if self.rolling_bars_storage[symbol][-1].Close >= self.EMA_ten_dict[symbol].Current.Value:

                                        # Quantity
                                        quantity = self.Portfolio[stock_puts[value]].Quantity

                                        # If quantity is 1 
                                        if quantity == 1:

                                            # Sell
                                            self.Liquidate(stock_puts[value], tag = "Put exit 1 - selling all holdings")

                                            # Pop
                                            stock_puts.pop(value)

                                        # Else
                                        else:

                                            # Total cost of position
                                            cost_of_position = average_price * quantity

                                            # Value of position 
                                            value_of_position = current_price * quantity

                                            # Profit
                                            profit = value_of_position - cost_of_position

                                            # If profit greater than 0
                                            if profit > 0:

                                                # Profit / current price
                                                quantity_to_keep = profit / current_price

                                                # Quantity to sell
                                                quantity_to_sell = int(quantity - quantity_to_keep)

                                                # If quantity_to_sell greater than or equal to 1
                                                if quantity_to_sell >= 1:

                                                    # Liquidate
                                                    self.MarketOrder(stock_puts[value], -quantity_to_sell, tag = "Put exit 1 - selling breakeven")                    
                    
                    # If length of call tradable list 2 greater than 0
                    if len(self.call_signal_two_tradable_list_two) > 0:

                        # Loop
                        for symbol in self.call_signal_two_tradable_list_two:

                            # Add to invested
                            invested_options.append(symbol)

                            # Call Entry 2
                            if (
                                
                                # If open greater than EMA 10
                                self.rolling_bars_storage[symbol][-1].Open > self.EMA_ten_dict[symbol].Current.Value and

                                # If RSI above 45
                                self.rolling_rsi_dict[symbol][-1] > 45 and

                                # If current RSI above previous RSI value
                                self.rolling_rsi_dict[symbol][-1] > self.rolling_rsi_dict[symbol][-2] and

                                # If current RSI above 2 previous RSI value
                                self.rolling_rsi_dict[symbol][-1] > self.rolling_rsi_dict[symbol][-3]
                                
                                ):

                                # Loop data
                                for kvp in self.data_storage:

                                    # Stock symbol
                                    stock_symbol = kvp.Value.Underlying.Symbol

                                    # If stock symbol in tradable list
                                    if stock_symbol == symbol:

                                        # Get option chain
                                        chain = kvp.Value

                                        # Sort call contracts with nearest strike and expiration
                                        contract = sorted(sorted(sorted(chain, \
                                            key = lambda x: abs(chain.Underlying.Price - x.Strike)), \
                                            key = lambda x: x.Expiry, reverse=False), \
                                            key = lambda x: x.Right, reverse=False)[0]

                                        # Buy 1 contract
                                        self.MarketOrder(contract.Symbol, 10, True, "Call entry 2")

                                        # Break
                                        break
                        
                        # Clear list
                        self.call_signal_two_tradable_list_two.clear()

                    # If length of put tradable list 2 greater than 0
                    if len(self.put_signal_two_tradable_list_two) > 0:

                        # Loop
                        for symbol in self.put_signal_two_tradable_list_two:

                            # Add to invested
                            invested_options.append(symbol)

                            # Put entry 2
                            if (

                                # If open less than EMA 10
                                self.rolling_bars_storage[symbol][-1].Open < self.EMA_ten_dict[symbol].Current.Value and

                                # If RSI below 55
                                self.rolling_rsi_dict[symbol][-1] < 55 and

                                # If current RSI below previous RSI value
                                self.rolling_rsi_dict[symbol][-1] < self.rolling_rsi_dict[symbol][-2] and

                                # If current RSI below 2 previous RSI value
                                self.rolling_rsi_dict[symbol][-1] < self.rolling_rsi_dict[symbol][-3]
                                
                                ):

                                # Loop data
                                for kvp in self.data_storage:

                                    # Stock symbol
                                    stock_symbol = kvp.Value.Underlying.Symbol

                                    # If stock symbol in tradable list
                                    if stock_symbol == symbol:

                                        # Get option chain
                                        chain = kvp.Value

                                        # Sort call contracts with nearest strike and expiration
                                        contract = sorted(sorted(sorted(chain, \
                                            key = lambda x: abs(chain.Underlying.Price - x.Strike)), \
                                            key = lambda x: x.Expiry, reverse=False), \
                                            key = lambda x: x.Right, reverse=True)[0]

                                        # Buy 1 contract
                                        self.MarketOrder(contract.Symbol, 10, True, "Put entry 2")

                                        # Break
                                        break
                
                        # Clear list
                        self.put_signal_two_tradable_list_two.clear()

                    # If length call tradable list 1 greater than 0
                    if len(self.call_signal_one_tradable_list) > 0:

                        # Loop data
                        for kvp in self.data_storage:

                            # Stock symbol
                            stock_symbol = kvp.Value.Underlying.Symbol

                            # If stock symbol in tradable list
                            if stock_symbol in self.call_signal_one_tradable_list:

                                # Add to invested
                                invested_options.append(symbol)

                                # If open of stock above EMA 10
                                if self.rolling_bars_storage[stock_symbol][-1].Open > self.EMA_ten_dict[stock_symbol].Current.Value:

                                    # Get option chain
                                    chain = kvp.Value

                                    # Sort call contracts with nearest strike and expiration
                                    contract = sorted(sorted(sorted(chain, \
                                        key = lambda x: abs(chain.Underlying.Price - x.Strike)), \
                                        key = lambda x: x.Expiry, reverse=False), \
                                        key = lambda x: x.Right, reverse=False)[0]

                                    # Buy 1 contract
                                    self.MarketOrder(contract.Symbol, 10, True, "Call entry 1")

                                    # Break
                                    break

                        # Clear
                        self.call_signal_one_tradable_list.clear()

                    # If length put tradable list 1 greater than 0
                    if len(self.put_signal_one_tradable_list) > 0:

                        # Loop data
                        for kvp in self.data_storage:

                            # Stock symbol
                            stock_symbol = kvp.Value.Underlying.Symbol

                            # If stock symbol in tradable list
                            if stock_symbol in self.put_signal_one_tradable_list:

                                # Add to invested
                                invested_options.append(symbol)

                                # If open of stock above EMA 10
                                if self.rolling_bars_storage[stock_symbol][-1].Open < self.EMA_ten_dict[stock_symbol].Current.Value:

                                    # Get option chain
                                    chain = kvp.Value

                                    # Sort put contracts with nearest strike and expiration
                                    contract = sorted(sorted(sorted(chain, \
                                        key = lambda x: abs(chain.Underlying.Price - x.Strike)), \
                                        key = lambda x: x.Expiry, reverse=False), \
                                        key = lambda x: x.Right, reverse=True)[0]

                                    # Buy 1 contract
                                    self.MarketOrder(contract.Symbol, 10, True, "Put entry 1")

                        # Clear
                        self.put_signal_one_tradable_list.clear()

                    # If length call tradable list 2 greater than 0
                    if len(self.call_signal_two_tradable_list_one) > 0:

                        # Loop
                        for symbol in self.call_signal_two_tradable_list_one:

                            # Add to invested
                            invested_options.append(symbol)

                            # Price less than or equal to EMA 10
                            if self.rolling_bars_storage[symbol][-1].Close <= self.EMA_ten_dict[symbol].Current.Value:

                                # Add to call signal 2 list 2
                                self.call_signal_two_tradable_list_two.append(symbol)

                        # Clear
                        self.call_signal_two_tradable_list_one.clear()

                    # If length put tradable list 2 greater than 0
                    if len(self.put_signal_two_tradable_list_one) > 0:

                        # Loop
                        for symbol in self.put_signal_two_tradable_list_one:

                            # Add to invested
                            invested_options.append(symbol)

                            # Price greater than or equal EMA 10
                            if self.rolling_bars_storage[symbol][-1].Close >= self.EMA_ten_dict[symbol].Current.Value:

                                # Add to put signal 2 list 2
                                self.put_signal_two_tradable_list_two.append(symbol)

                        # Clear
                        self.put_signal_two_tradable_list_one.clear()

                    # Loop through stocks
                    for symbol in self.symbols_list:

                        # If stock not invested
                        if symbol not in invested_options:

                            # Call entry logic 1
                            if (
                                
                                # Previous EMA difference is negative
                                self.previous_ema_one_ema_three_difference[symbol][-2] < 0 and 
                                
                                # If current EMA diffeerence is positive
                                self.previous_ema_one_ema_three_difference[symbol][-1] > 0 and

                                # If EMA 1 is above EMA 4
                                self.EMA_ten_dict[symbol].Current.Value > self.EMA_forty_dict[symbol].Current.Value and

                                # If current candle close above EMA 1
                                self.rolling_bars_storage[symbol][-1].Close > self.EMA_ten_dict[symbol].Current.Value and

                                # If RSI above 45
                                self.rolling_rsi_dict[symbol][-1] > 45

                                # If current RSI above previous RSI value
                                and self.rolling_rsi_dict[symbol][-1]  > self.rolling_rsi_dict[symbol][-2]

                                # If current RSI above 2 previous RSI value
                                and self.rolling_rsi_dict[symbol][-1]  > self.rolling_rsi_dict[symbol][-3]
                                
                                ):

                                # Add to call tradable list 1
                                # self.call_signal_one_tradable_list.append(symbol)

                                self.MarketOrder(symbol, 1, tag = " Call 1  ")

                            # Call entry logic 2
                            elif (
                                
                                # If EMA 1 above EMA 3
                                self.EMA_ten_dict[symbol].Current.Value > self.EMA_thirty_dict[symbol].Current.Value and

                                # If EMA 1 above EMA 4
                                self.EMA_ten_dict[symbol].Current.Value > self.EMA_forty_dict[symbol].Current.Value and

                                # If EMA 2 above EMA 4
                                self.EMA_twenty_dict[symbol].Current.Value > self.EMA_forty_dict[symbol].Current.Value and

                                # If EMA 3 above EMA 4
                                self.EMA_thirty_dict[symbol].Current.Value > self.EMA_forty_dict[symbol].Current.Value and

                                # If price above EMA 1
                                self.rolling_bars_storage[symbol][-1].Close > self.EMA_ten_dict[symbol].Current.Value
                                
                                ):

                                # Add to call tradable list 2
                                # self.call_signal_two_tradable_list_one.append(symbol)
                                self.MarketOrder(symbol, 1, tag = " Call 2  ")

                            # Put entry logic 1
                            elif (
                                
                                # If previous EMA difference positive
                                self.previous_ema_one_ema_three_difference[symbol][-2] > 0 and 
                                
                                # If current EMA diffeerence is negative
                                self.previous_ema_one_ema_three_difference[symbol][-1] < 0 and

                                # If EMA 1 is below EMA 4
                                self.EMA_ten_dict[symbol].Current.Value < self.EMA_forty_dict[symbol].Current.Value and

                                # If current candle close below EMA 1
                                self.rolling_bars_storage[symbol][-1].Close < self.EMA_ten_dict[symbol].Current.Value and

                                # If RSI below 55
                                self.rolling_rsi_dict[symbol][-1] < 55 and

                                # If current RSI below previous RSI value
                                self.rolling_rsi_dict[symbol][-1]  < self.rolling_rsi_dict[symbol][-2] and

                                # If current RSI below 2 previous RSI value
                                self.rolling_rsi_dict[symbol][-1]  < self.rolling_rsi_dict[symbol][-3] 
                                
                                ):

                                # Add to put signal 1 list
                                # self.put_signal_one_tradable_list.append(symbol)
                                self.MarketOrder(symbol, 1, tag = " Put 1  ")

                            # Put entry logic 2
                            elif (
                                
                                # If EMA 1 below EMA 3
                                self.EMA_ten_dict[symbol].Current.Value < self.EMA_thirty_dict[symbol].Current.Value and

                                # If EMA 1 below EMA 4
                                self.EMA_ten_dict[symbol].Current.Value < self.EMA_forty_dict[symbol].Current.Value and

                                # If EMA 2 below EMA 4
                                self.EMA_twenty_dict[symbol].Current.Value < self.EMA_forty_dict[symbol].Current.Value and

                                # If EMA 3 below EMA 4
                                self.EMA_thirty_dict[symbol].Current.Value < self.EMA_forty_dict[symbol].Current.Value and

                                # If price below EMA 1
                                self.rolling_bars_storage[symbol][-1].Close < self.EMA_ten_dict[symbol].Current.Value
                                
                                ):

                                # Add to put tradable list 2
                                # self.put_signal_two_tradable_list_one.append(symbol)
                                self.MarketOrder(symbol, 1, tag = " Put 2  ")

                # If current time is exactly 2:45 pm
                elif current_time == two_forty_five:

                    # Sell everything
                    self.Liquidate(tag = "Liquidated at 2:45 pm")

                    # Reset
                    self.call_signal_one_tradable_list.clear()

                    # Call tradable list for signal 2 - 1
                    self.call_signal_two_tradable_list_one.clear()

                    # Call tradable list for signal 2 - 2
                    self.call_signal_two_tradable_list_two.clear()

                    # Put tradable list for signal 1
                    self.put_signal_one_tradable_list.clear()

                    # Put tradable list for signal 2 - 1
                    self.put_signal_two_tradable_list_one.clear()

                    # Put tradable list for signal 2 - 2
                    self.put_signal_two_tradable_list_two.clear()