# region imports
from AlgorithmImports import *
# endregion

# Import statistics
import statistics

# Import datetime
import datetime

class VirtualGreenHornet(QCAlgorithm):

    def Initialize(self):

        # Set Start Date
        self.SetStartDate(2022, 10, 1)

        # Set End Date
        # self.SetEndDate(2022, 10, 5)

        # Set Strategy Cash
        self.SetCash(100000)

        # Add universe
        self.AddUniverse(self.Coarse, self.Fine)

        # Add SPY
        self.SPY_symbol = self.AddEquity("SPY", Resolution.Minute, extendedMarketHours = True).Symbol
        
        # Universe extended hours
        self.UniverseSettings.ExtendedMarketHours = True

        # Universe resolution
        self.UniverseSettings.Resolution = Resolution.Minute

        # # # Storage # # #

        # Symbols storage
        self.symbols_list = []

        # Consolidator dictionary
        self.consolidator_dictionary = {}

        # 9:30 am price
        self.day_starting_price = {}

        # Volume
        self.day_candles = {}

        # List to store stocks that pass scanning component
        self.list_1 = []

        # Counter
        self.counter = 0

        # 5-minute consolidator dictionary
        self.five_minute_consolidator_dictionary = {}

        # 5-minute bar dictionary
        self.five_minute_bar_dictionary = {}

        # Counter 2
        self.counter_2 = 0

        # 10-minute consolidator dictionary
        self.ten_minute_consolidator_dictionary = {}

        # 5-minute bar dictionary
        self.ten_minute_bar_dictionary = {}

        # Counter 3
        self.counter_3 = 0

        # List 2
        self.list_2 = []

        # RSI storage
        self.RSI_storage = {}

        # RSI values storage
        self.RSI_values = {}

        # Warm up
        self.warm_up = False

        # # # End storage # # #

        # Schedule 1
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                    self.TimeRules.AfterMarketOpen("SPY", 1),
                    self.MinuteAfterMarketOpen)

        # Schedule 2
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                    self.TimeRules.BeforeMarketClose("SPY", 5),
                    self.LiquidateHoldings)

    def Coarse(self, coarse):

        # Return list
        return_list = []

        # # List
        list_one = [x.Symbol for x in coarse if x.HasFundamentalData and x.Price > 0.4]

        # # History
        history = self.History(list_one, 90, Resolution.Daily)

        # # Loop
        for symbol in list_one:

            # Check if has history
            if str(symbol) in history.index:

                # Loc
                symbol_history = history.loc[symbol]

                # If length equal to 90
                if len(symbol_history) == 90:

                    # Average volume
                    average_volume = symbol_history["volume"].mean()

                    # If average volume above 500,000
                    if average_volume > 500000:

                        # Add to return list
                        return_list.append(symbol)

        # Return
        return return_list

    def Fine(self, fine):

        # Return
        return [x.Symbol for x in fine if x.CompanyReference.PrimaryExchangeID == "NAS"]

    def OnSecuritiesChanged(self, changes):

        # Symbols 
        symbols = [x.Symbol for x in changes.AddedSecurities]

        # 6-period rsi
        random_rsi = RelativeStrengthIndex(6)

        # Warm up period
        warm_up_period = random_rsi.WarmUpPeriod * 5

        # History
        history = self.History(symbols, warm_up_period, Resolution.Minute)

        # Warm up set to False
        self.warm_up = False

        problem = []

        # Loop added
        for symbol in symbols:

            # Check if symbol is SPY
            if symbol != self.SPY_symbol:

                # If symbol has historical data
                if str(symbol) in history.index:

                    # Loc
                    symbol_history = history.loc[symbol]

                    # If length == requested
                    if len(symbol_history) == warm_up_period:

                        # Add to symbols list
                        self.symbols_list.append(symbol)
                
                        # Consolidator
                        self.consolidator_dictionary[symbol] = TradeBarConsolidator(timedelta(minutes = 15))

                        # Consolidated handler
                        self.consolidator_dictionary[symbol].DataConsolidated += self.consolidation_handler

                        # Register the consolidator for automatic updates
                        self.SubscriptionManager.AddConsolidator(symbol, self.consolidator_dictionary[symbol])

                        # Consolidator
                        self.five_minute_consolidator_dictionary[symbol] = TradeBarConsolidator(timedelta(minutes = 5))

                        # Consolidated handler
                        self.five_minute_consolidator_dictionary[symbol].DataConsolidated += self.five_minute_consolidation_handler

                        # Register the consolidator for automatic updates
                        self.SubscriptionManager.AddConsolidator(symbol, self.five_minute_consolidator_dictionary[symbol])

                        # 10-minute consolidator
                        self.ten_minute_consolidator_dictionary[symbol] = TradeBarConsolidator(timedelta(minutes = 10))

                        # Consolidated handler
                        self.ten_minute_consolidator_dictionary[symbol].DataConsolidated += self.ten_minute_consolidation_handler

                        # Register the consolidator for automatic updates
                        self.SubscriptionManager.AddConsolidator(symbol, self.ten_minute_consolidator_dictionary[symbol])

                        # RSI
                        self.RSI_storage[symbol] = RelativeStrengthIndex(6)

                        # Register
                        self.RegisterIndicator(symbol = symbol, indicator = self.RSI_storage[symbol], consolidator = self.five_minute_consolidator_dictionary[symbol])

                        # Loop
                        for time, row in symbol_history.iterrows():
                            
                            # Try
                            try:

                                # Create candlestick
                                bar = TradeBar(time - datetime.timedelta(minutes = 1), symbol, row.open, row.high, row.low, row.close, row.volume, datetime.timedelta(minutes = 1))

                                # Update
                                self.five_minute_consolidator_dictionary[symbol].Update(bar)

                            # Except
                            except:
                                
                                # Symbol
                                if symbol not in problem:
                                    
                                    # Add
                                    problem.append(symbol)

        # Warm up set to True
        self.warm_up = True

        # Loop
        for symbol in problem:

            # Unregister consolidator
            self.SubscriptionManager.RemoveConsolidator(symbol, self.consolidator_dictionary[symbol])

            # Remove consolidator
            self.consolidator_dictionary.pop(symbol)

            # Remove
            self.symbols_list.remove(symbol)

            # Unregister consolidator
            self.SubscriptionManager.RemoveConsolidator(symbol, self.five_minute_consolidator_dictionary[symbol])

            # Remove consolidator
            self.five_minute_consolidator_dictionary.pop(symbol)

            # Unregister consolidator
            self.SubscriptionManager.RemoveConsolidator(symbol, self.ten_minute_consolidator_dictionary[symbol])

            # Remove consolidator
            self.ten_minute_consolidator_dictionary.pop(symbol)

            # RSI
            self.RSI_storage.pop(symbol)

        # Loop removed
        for security in changes.RemovedSecurities:
            
            # Symbol
            symbol = security.Symbol

            # If symbol in symbols list
            if symbol in self.symbols_list:

                # Unregister consolidator
                self.SubscriptionManager.RemoveConsolidator(symbol, self.consolidator_dictionary[symbol])

                # Remove consolidator
                self.consolidator_dictionary.pop(symbol)

                # Remove
                self.symbols_list.remove(symbol)

                # Unregister consolidator
                self.SubscriptionManager.RemoveConsolidator(symbol, self.five_minute_consolidator_dictionary[symbol])

                # Remove consolidator
                self.five_minute_consolidator_dictionary.pop(symbol)

                # Unregister consolidator
                self.SubscriptionManager.RemoveConsolidator(symbol, self.ten_minute_consolidator_dictionary[symbol])

                # Remove consolidator
                self.ten_minute_consolidator_dictionary.pop(symbol)

                # RSI
                self.RSI_storage.pop(symbol)

    def MinuteAfterMarketOpen(self):

        # History
        history = self.History(self.symbols_list, 1, Resolution.Daily)

        # Loop
        for symbol in self.symbols_list:

            # If symbol has data
            if str(symbol) in history.index:

                # If price not more than 60% from yesterday close
                if self.Securities[symbol].Open < history.loc[symbol]["close"][-1] * 1.6:

                    # Add to day starting price
                    self.day_starting_price[symbol] = self.Securities[symbol].Open

                    # Add to day candles
                    self.day_candles[symbol] = []

                    # Add to 5-minute bar storage
                    self.five_minute_bar_dictionary[symbol] = []

                    # Add to 10-minute bar storage
                    self.ten_minute_bar_dictionary[symbol] = []

                    # Add to RSI values
                    self.RSI_values[symbol] = []

    def consolidation_handler(self, sender, bar):

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

        # 11 am
        eleven = datetime.time(11, 00)

        # 3:55 pm
        three = datetime.time(15, 55)

        # If symbol in day volume 
        if bar.Symbol in self.day_candles:

            # Add
            self.day_candles[bar.Symbol].append(bar)
        
        # Add to counter
        self.counter += 1

        # If counter equal to number of stocks in algo
        if self.counter == len(self.consolidator_dictionary):

            # Reset
            self.counter = 0

            # If time is 11
            if current_time >= eleven and current_time <= three:

                # History symbols list
                history_symbols_list = []

                # Loop
                for symbol in self.day_starting_price:
                    
                    # If price greater than 5% from opening price
                    if self.Securities[symbol].Close > self.day_starting_price[symbol] * 1.05:

                        # Add
                        history_symbols_list.append(symbol)

                # If length
                if len(history_symbols_list) > 0:
                
                    # Get history
                    history = self.History(history_symbols_list[-1], 5, Resolution.Daily)

                    # Get timestamp
                    timestamp = history.index[-1][-1] - timedelta(days = 1)

                    # Start datetime
                    start_datetime = datetime.datetime(timestamp.year, timestamp.month, timestamp.day, 9, 30)

                    # End datetime
                    end_datetime = datetime.datetime(timestamp.year, timestamp.month, timestamp.day, self.Time.hour, self.Time.minute)
                        
                    # Get history
                    history = self.History(history_symbols_list, start_datetime, end_datetime, Resolution.Minute)

                    # Loop
                    for symbol in history_symbols_list:

                        # If symbol in history
                        if str(symbol) in history.index:

                            # Get current volume
                            current_volume = sum([x.Volume for x in self.day_candles[symbol]])

                            # If current volume greater than 300,000
                            if current_volume > 300000:

                                # Symbol history
                                symbol_history = history.loc[symbol]

                                # Previous day volume
                                previous_volume = symbol_history["volume"].sum()

                                # If current volume 5 times previous volume
                                if current_volume > previous_volume * 5:

                                    # If not in list 1
                                    if symbol not in self.list_1:

                                        # Add
                                        self.list_1.append(symbol)

    def five_minute_consolidation_handler(self, sender, bar):

        # If not warm up
        if self.warm_up == True:

            # Add 
            self.counter_2 += 1

            # If symbol in 5-minute bar dictionary
            if bar.Symbol in self.five_minute_bar_dictionary:

                # Add
                self.five_minute_bar_dictionary[bar.Symbol].append(bar)

            # If symbol in RSI values storage
            if bar.Symbol in self.RSI_values:

                # Add
                self.RSI_values[bar.Symbol].append(self.RSI_storage[bar.Symbol].Current.Value)

            # If counter equal length
            if self.counter_2 == len(self.consolidator_dictionary):

                # Reset
                self.counter_2 = 0

                # If length of list 2 above 0
                if len(self.list_2) > 0:

                    # Loop
                    for symbol in self.list_2:

                        # If not invested
                        if not self.Portfolio[symbol].Invested:

                            # Past 3 RSI
                            past_3_RSI = self.RSI_values[symbol][-3:]

                            # If max above 70
                            if max(past_3_RSI) > 70:

                                # If current bar is red
                                if self.five_minute_bar_dictionary[symbol][-1].Close < self.five_minute_bar_dictionary[symbol][-1].Open:

                                    # If 2 bars prior is green
                                    if self.five_minute_bar_dictionary[symbol][-3].Open < self.five_minute_bar_dictionary[symbol][-3].Close:

                                        # If current bar bigger than previous
                                        if abs(self.five_minute_bar_dictionary[symbol][-1].Close - self.five_minute_bar_dictionary[symbol][-1].Open) > abs(self.five_minute_bar_dictionary[symbol][-2].Close - self.five_minute_bar_dictionary[symbol][-2].Open):

                                            # If 2 bars prior bigger than previous
                                            if abs(self.five_minute_bar_dictionary[symbol][-3].Close - self.five_minute_bar_dictionary[symbol][-3].Open) > abs(self.five_minute_bar_dictionary[symbol][-2].Close - self.five_minute_bar_dictionary[symbol][-2].Open):

                                                # Short
                                                self.MarketOrder(symbol, -1)

    def ten_minute_consolidation_handler(self, sender, bar):

        # Add 
        self.counter_3 += 1

        # If symbol in 10-minute bar dictionary
        if bar.Symbol in self.ten_minute_bar_dictionary:

            # Add
            self.ten_minute_bar_dictionary[bar.Symbol].append(bar)

        # If counter equal length
        if self.counter_3 == len(self.consolidator_dictionary):

            # Reset
            self.counter_3 = 0

            # If length of list 1 above 0
            if len(self.list_1) > 0:

                # Loop
                for symbol in self.list_1:

                    # Bars
                    bars = self.ten_minute_bar_dictionary[symbol]

                    # Recent bar value
                    recent_bar_value = bars[-1].Close - bars[-1].Open

                    # List
                    new_list = [(x.Close - x.Open) for x in bars]

                    # Reverse
                    new_list.reverse()

                    # Bars reverse
                    bars.reverse()

                    # If red
                    if recent_bar_value < 0:

                        # Loop
                        for value in new_list:

                            # If value greater than 0
                            if value > 0:

                                # Get value as direcional change
                                directional_change = bars[new_list.index(value)].Close

                                # Get highest close
                                highest_close = max([x.Close for x in bars])

                                # Get open
                                open_price = self.day_starting_price[symbol]

                                # Difference
                                difference = highest_close - open_price

                                # 40% of difference
                                fourty_percent = difference * 0.4

                                # If current close - directional change more than 40% of difference
                                if self.ten_minute_bar_dictionary[symbol][-1].Close - directional_change > fourty_percent:

                                    # Add to list 2
                                    self.list_2.append(symbol)

                    # If green
                    elif recent_bar_value > 0:

                        # Loop
                        for value in new_list:

                            # If value less than 0
                            if value < 0:

                                # Get value as direcional change
                                directional_change = bars[new_list.index(value)].Close

                                # Get highest close
                                highest_close = max([x.Close for x in bars])

                                # Get open
                                open_price = self.day_starting_price[symbol]

                                # Difference
                                difference = highest_close - open_price

                                # 40% of difference
                                fourty_percent = difference * 0.4

                                # If current close - directional change more than 40% of difference
                                if self.ten_minute_bar_dictionary[symbol][-1].Close - directional_change > fourty_percent:

                                    # Add to list 2
                                    self.list_2.append(symbol)

    def LiquidateHoldings(self):

        # Sell everything
        self.Liquidate()
        
        # Remove
        self.list_1.clear()

        # Remove
        self.day_starting_price.clear()

        # Remove
        self.day_candles.clear()

        # Remove
        self.five_minute_bar_dictionary.clear()

        # Remove
        self.ten_minute_bar_dictionary.clear()

        # Remove
        self.list_2.clear()

        # Remove
        self.RSI_values.clear()

    def OnData(self, data: Slice):
        pass