from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

class QuantumHorizontalRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 3, 24)  # Set Start Date
        #self.SetEndDate(2019, 11, 14)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        self.AddEquity("SPY", Resolution.Minute)
        self.scaning = False
        self.lastToggle = None
        
        self.__numberOfSymbols = 100
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))

        self.AddAlpha(FadeTheGapModel(self))
        
        self.SetExecution(ImmediateExecutionModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("SPY", 0), self.toggleScan)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen("SPY", 45), self.toggleScan)
        
    def toggleScan(self):
        self.scaning = not self.scaning
        self.lastToggle = self.Time
        

    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 FadeTheGapModel(AlphaModel):
    yest_closes = {}
    
    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['SPY'].Exchange.ExchangeOpen:
            for symbol, yest_close in self.yest_closes.items():
                if symbol in slice.Bars:
                    self.yest_closes[symbol] = slice.Bars[symbol].Close
    
        if not self.algo.scaning:
            return []
        
        insights = []
        
        duration = algorithm.Time - self.algo.lastToggle
        insight_mins = 45 - (duration.seconds // 60 % 60)
        
        # Create insights for symbols up atleast 10% on the day
        for symbol, yest_close in self.yest_closes.items():
            # If already invested, continue to next symbol
            if algorithm.Securities[symbol].Invested:
                continue
            
            # Calculate return sign yesterday's close
            ret = (slice[symbol].Close - yest_close) / yest_close
            if ret >= 0.1: # Up 10% on the day
                insights.append(Insight(symbol, timedelta(minutes=insight_mins), InsightType.Price, InsightDirection.Down))
        
        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)['close']
            
            for added in changes.AddedSecurities:
                # Save yesterday's close
                closes = history.loc[[str(added.Symbol.ID)]].values
                if len(closes) < 1:
                    continue
                self.yest_closes[added.Symbol] = closes[0]
            
        for removed in changes.RemovedSecurities:
            # Delete yesterday's close tracker
            self.yest_closes.pop(removed.Symbol, None)