class QuantumHorizontalRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 15)  # Set Start Date
        self.SetEndDate(2021, 1, 31)
        self.SetCash(333333)  # Set Strategy Cash
        self.AddEquity("W5000", Resolution.Minute)
        self.scaning = False
        self.lastToggle = None
        self.needs_reset = False
        
        self.__numberOfSymbols =10
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))
        self.UniverseSettings.Resolution = Resolution.Minute

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

        self.SetPortfolioConstruction(AccumulativeInsightPortfolioConstructionModel(lambda time: None))
        
        self.SetRiskManagement(MaximumUnrealizedProfitPercentPerSecurity(maximumUnrealizedProfitPercent = 0.02))  

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("W5000", 0), self.toggleScan)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("W5000", 375), self.toggleScan)
        
    def toggleScan(self):
        self.scaning = not self.scaning
        self.lastToggle = self.Time
        
        if not self.scaning:
            self.needs_reset = True
            

    def CoarseSelectionFunction(self, coarse):
        # Stocks with the most dollar volume traded yesterday
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]
    
    def FineSelectionFunction(self, fine):
        return [ x.Symbol for x in fine ]
        

class ShortSqueezeModel(AlphaModel):
    symbolData = {}

    def __init__(self, algo):
        self.algo = algo
        
    
    def Update(self, algorithm, slice):
        if algorithm.IsWarmingUp:
            return []
    
        # If it's the end of the day, update the yesterday close of each indicator
        if not algorithm.Securities['W5000'].Exchange.ExchangeOpen:
            for symbol in self.symbolData:
                if symbol in slice.Bars:
                    self.symbolData[symbol].yest_close = slice.Bars[symbol].Close
    
        if not self.algo.scaning:
            # Reset max indicator
            if self.algo.needs_reset:
                for symbol in self.symbolData:
                    self.symbolData[symbol].max.Reset()
                self.algo.needs_reset = False
            return []
        
        insights = []
        
        insight_seconds = 99999999999
        
        # Create insights for symbols up at least 10% on the day
        for symbol in self.symbolData:
            # If already invested, continue to next symbol
            if algorithm.Securities[symbol].Invested or symbol not in slice.Bars or self.symbolData[symbol].max.Samples == 0:
                continue
            
            # Calculate return sign yesterday's close
            yest_close = self.symbolData[symbol].yest_close
            close = slice[symbol].Close 
            ret = (close - yest_close) / yest_close
            high_of_day_break = close > self.symbolData[symbol].max.Current.Value
            if ret >= 0.4 and high_of_day_break:    # Up 40% on the day & breaks high of day
                hours = algorithm.Securities[symbol].Exchange.Hours
        # 1-minute before the close
                closeTime = hours.GetNextMarketClose(algorithm.Time, False) - timedelta(minutes=5)
                insights.append(Insight.Price(symbol, closeTime, InsightDirection.Up))
                
                
        # Update max indicator for all symbols
        for symbol in self.symbolData:
            if symbol in slice.Bars:
                self.symbolData[symbol].max.Update(slice.Time, slice.Bars[symbol].High)
                
                
        # Constantly updating 5% Trailing Stop Order        
        for symbol in self.symbolData:
            if symbol in slice.Bars and algorithm.Securities[symbol].Invested and slice[symbol].Close <= 0.98*self.symbolData[symbol].max.Current.Value:
                insights.append(Insight(symbol, timedelta(seconds=insight_seconds), InsightType.Price, InsightDirection.Flat))
        
        return Insight.Group(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])
            
        for removed in changes.RemovedSecurities:
            # Delete yesterday's close tracker
            self.symbolData.pop(removed.Symbol, None)

class SymbolData:
    def __init__(self, yest_close):
        self.yest_close = yest_close
        self.max = Maximum(45*60) # 45 minutes
        
        
        
class MaximumUnrealizedProfitPercentPerSecurity(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the unrealized profit per holding to the specified percentage'''

    def __init__(self, maximumUnrealizedProfitPercent = 0.74):
        '''Initializes a new instance of the MaximumUnrealizedProfitPercentPerSecurity class
        Args:
            maximumUnrealizedProfitPercent: The maximum percentage unrealized profit allowed for any single security holding, defaults to 5% drawdown per security'''
        self.maximumUnrealizedProfitPercent = abs(maximumUnrealizedProfitPercent)
        
    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
        targets = []
        for kvp in algorithm.Securities:
            security = kvp.Value

            if not security.Invested:
                continue

            pnl = security.Holdings.UnrealizedProfitPercent
            if pnl > self.maximumUnrealizedProfitPercent:

                ### For debugging, add this to see when it is being called
                algorithm.Log('Risk model triggered')

                # liquidate
                targets.append(PortfolioTarget(security.Symbol, 0))

        return targets