# region imports
from AlgorithmImports import *
import time
# endregion

class FetchTopGappers(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 5, 11)  # Set Start Date
        self.SetEndDate(2021, 5, 11)
        self.SetCash(100000)  # Set Strategy Cash

        self.AddEquity("SPY", resolution=Resolution.Minute, extendedMarketHours=True)        

        # Add universe
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.ExtendedMarketHours = True
        self.UniverseSettings.FillForward = True
        self.AddUniverse(self.CoarseUniverseSelection)

        self.AddUniverseSelection(ScheduledUniverseSelectionModel(
            self.DateRules.EveryDay("SPY"), 
            self.TimeRules.AfterMarketOpen("SPY", -5),
            self.ScheduledSymbolSelect))

        self.cur_time = time.time()
        self.num_gappers = 5

    def OnData(self, data: Slice):
        if (self.Time.hour == 8 and self.Time.minute == 31):
            self.Debug(f"Time: {self.Time} | Length of slice: {len(data.Bars.Values)}")
        
        if (self.Time.hour == 9 and self.Time.minute == 31):
            self.Debug(f"Time: {self.Time} | Length of slice: {len(data.Bars.Values)}")

    def CoarseUniverseSelection(self, coarse):
        self.all_relevant_stocks = {}
        self.Debug(f"In coarse at {self.Time}, time: {int(time.time() - self.cur_time)} seconds")
        self.cur_time = time.time()

        for stock in coarse:
            if (stock.Symbol.Value == "SPY"):
                continue
            if (stock.AdjustedPrice > 1):
                self.all_relevant_stocks[str(stock.Symbol)] = {"symbol": stock.Symbol, "last_close": stock.Price}   
        return []
    
    def calc_gap(self, row):
        symbol_data = self.all_relevant_stocks[row.name]
        last_close = symbol_data["last_close"]
        cur_price = row['close']
        gap = ((cur_price * 100) / last_close) - 100
        return {"ticker": symbol_data["symbol"].Value, "gap": gap, "symbol": symbol_data["symbol"], "last_close": last_close, "cur_price": cur_price}
    
    def ScheduledSymbolSelect(self, date):
        min_history = self.History(
            tickers=[x["symbol"] for x in self.all_relevant_stocks.values()], 
            start=self.Time - timedelta(minutes=1),
            end=self.Time,
            resolution=Resolution.Minute, 
            fillForward=True, 
            extendedMarket=True,
            dataNormalizationMode=DataNormalizationMode.Adjusted, 
        ) 

        if (min_history.empty):
            return []
        
        # Drop timestamp level
        min_history = min_history.droplevel(1)
        # Calculate gappers
        gappers = list(min_history.apply(self.calc_gap, axis=1))
        gappers = sorted(gappers, key=lambda d: d['gap'], reverse=True)
        self.top_gappers = gappers[:self.num_gappers]

        self.Debug(f"{self.top_gappers[0]}")
        self.Debug(f"{self.top_gappers[1]}")
        self.selected_symbols = [stock["symbol"] for stock in self.top_gappers]
        return self.selected_symbols