class QuantumHorizontalRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 7, 1)  # Set Start Date
        self.SetEndDate(2020, 8, 31)
        self.SetCash(333333)  # Set Strategy Cash
        
        
        self.backtestSymbolsPerDay = {}
        self.current_universe = []

        self.UniverseSettings.Resolution = Resolution.Second
        
        self.SetUniverseSelection(ScheduledUniverseSelectionModel(
            self.DateRules.Every(DayOfWeek.Monday, DayOfWeek.Tuesday, DayOfWeek.Wednesday, DayOfWeek.Thursday, DayOfWeek.Friday),
            self.TimeRules.Every(timedelta(minutes = 15)),
            self.SelectSymbols
            ))

        self.AddAlpha(ShortSqueezeModel(self))
        
        self.SetExecution(ImmediateExecutionModel())

        self.SetPortfolioConstruction(AccumulativeInsightPortfolioConstructionModel(lambda time: None))
        
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.02))
        

    def SelectSymbols(self, dateTime):
         # handle live mode file format
        if self.LiveMode:
            # fetch the file from dropbox
            str = self.Download("https://www.dropbox.com/s/x4ptyv396uv0ezu/Watchlist%202.csv?dl=1")
            # if we have a file for today, return symbols, else leave universe unchanged
            self.current_universe = str.split(',') if len(str) > 0 else self.current_universe
            return self.current_universe

        # backtest - first cache the entire file
        if len(self.backtestSymbolsPerDay) == 0:

            # No need for headers for authorization with dropbox, these two lines are for example purposes 
            #byteKey = base64.b64encode("UserName:Password".encode('ASCII'))
            # The headers must be passed to the Download method as dictionary
            #headers = { 'Authorization' : f'Basic ({byteKey.decode("ASCII")})' }

            #str = self.Download("https://www.dropbox.com/s/x4ptyv396uv0ezu/Watchlist%202.csv?dl=1", headers)
            csv_data = self.Download("https://www.dropbox.com/s/x4ptyv396uv0ezu/Watchlist%202.csv?dl=1")
            for line in csv_data.splitlines():
                data = line.split(',')
                self.backtestSymbolsPerDay[data[0]] = data[1:]
        
        index = dateTime.date().strftime("%Y%m%d")
        
        if index in self.backtestSymbolsPerDay.keys():
            tickers = self.backtestSymbolsPerDay[index]
            self.current_universe = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in tickers]
        
        return self.current_universe
        

class ShortSqueezeModel(AlphaModel):
    symbolData = {}

    def __init__(self, algo):
        self.algo = algo
        
    
    def Update(self, algorithm, slice):
        insights = []
        
        # Create insights for symbols up at least 40% on the day
        for symbol, symbol_data in self.symbolData.items():
            
            # If already invested, continue to next symbol
            if algorithm.Securities[symbol].Invested or symbol not in slice.Bars or not symbol_data.max.IsReady:
                continue
            
            # Calculate return sign yesterday's close
            yest_close = symbol_data.yest_close
            close = slice[symbol].Close 
            ret = (close - yest_close) / yest_close
            high_of_day_break = close > symbol_data.max.Current.Value
            
            if ret >= 0.4 and high_of_day_break:    # Up 20% on the day & breaks high of day
                hours = algorithm.Securities[symbol].Exchange.Hours 
                closeTime = hours.GetNextMarketClose(algorithm.Time, False) - timedelta(minutes=2) # exit 1-minute before the close
                insights.append(Insight.Price(symbol, closeTime, InsightDirection.Up))
        
        return insights
    
    
    
    def OnSecuritiesChanged(self, algorithm, changes):
        if len(changes.AddedSecurities) > 0:
            # Get history of symbols over lookback window
            added_symbols = [x.Symbol for x in changes.AddedSecurities]
            history = algorithm.History(added_symbols, 1, Resolution.Daily)
            if history.empty:
                    return
            history  = history['close']
            for added in changes.AddedSecurities:
                # Save yesterday's close
                closes = history.loc[[str(added.Symbol.ID)]].values
                if len(closes) < 1:
                    continue
                self.symbolData[added.Symbol] = SymbolData(closes[0], algorithm, added.Symbol)
            
        for removed in changes.RemovedSecurities:
            # Delete yesterday's close tracker
            symbol_data = self.symbolData.pop(removed.Symbol, None)
            if symbol_data:
                symbol_data.dispose()

class SymbolData:
    def __init__(self, yest_close, algo, symbol):
        self.yest_close = yest_close
        self.algo = algo
        self.symbol = symbol
        self.event = algo.Schedule.On(algo.DateRules.EveryDay(symbol), algo.TimeRules.AfterMarketOpen(symbol, 375), self.reset)
        self.internal_max = Maximum(45) # 45 minutes
        self.max = Delay(1) 
        self.daily_consolidator = TradeBarConsolidator(timedelta(days = 1))    ## 1 Day TradeBar Consolidator
        self.daily_consolidator.DataConsolidated += self.DailyConsolidator  ## Add fuction to do what you want every day with your data
        self.algo.SubscriptionManager.AddConsolidator(self.symbol, self.daily_consolidator)
        
        self.minute_consolidator = TradeBarConsolidator(timedelta(minutes = 1))    ## 1 Day TradeBar Consolidator
        self.minute_consolidator.DataConsolidated += self.MinuteConsolidator  ## Add fuction to do what you want every day with your data
        self.algo.SubscriptionManager.AddConsolidator(self.symbol, self.minute_consolidator)
    
    def MinuteConsolidator(self, sender, bar):
        self.internal_max.Update(bar.Time,bar.High)
        self.max.Update(bar.Time, self.internal_max.Current.Value)
        
    def DailyConsolidator(self, sender, bar):
        self.yest_close = bar.Close
        
    def reset(self):
        self.internal_max.Reset()
        self.max.Reset()
        
    def dispose(self):
        self.algo.Schedule.Remove(self.event)
        
        # Remove consolidators
        self.algo.SubscriptionManager.RemoveConsolidator(self.symbol, self.daily_consolidator)
        self.algo.SubscriptionManager.RemoveConsolidator(self.symbol, self.minute_consolidator)