from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
class TransdimensionalParticleThrustAssembly(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 5, 28)                                                  # Set Start Date
        self.SetEndDate(2020, 6, 5)                                                     # Set End Date
        self.SetCash(100000)                                                            # Set Strategy Cash
        self.AddEquity("SPY", Resolution.Minute).SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted) # Add SPY to set scheduled events
        self.UniverseSettings.Resolution = Resolution.Minute                            # Setting Universe: Daily, Minute or Second
        self.UniverseSettings.FillForward = True
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 2), self.Rebalance) # Our Scheduled Events
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 30), self.LiquidatePositions)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 1), self.OnMarketClose)
        self.previousClose = {}                                                         # Dictionary to keep track of previous close for each symbol
        self.Spliteventbefore = {}
        self.Spliteventafter = {}
        self.donottrade = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in []] #['HUGE']]#, 'MSFT']]
        self.cashused = 10000
        
    def OnData(self, data):                                                             # OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
        pass
    
    def CoarseSelectionFunction(self, coarse):                                          # Picks up securities Universe.  Constructed at midnight of night before.
        return [x.Symbol for x in coarse if 20 > x.Price]
    
    def FineSelectionFunction(self, fine):                                              # Picks up securities from Coarse > Universe.  Constructed at midnight of night before.
        return [x.Symbol for x in fine if x.MarketCap < 500000000]
        
    def OnSecuritiesChanged(self, changes):                                             # Picks up securities from the Fine > Coarse > Universe.  Constructed at midnight of night before.
        for security in changes.AddedSecurities:                                        # AddedSecurities are those populated by Fine > Coarse > Universe, for security in self.ActiveSecurities.Values
            if security.Symbol in self.donottrade:
                continue
            symbol = security.Symbol
            ## self.Spliteventbefore[symbol] = self.Value.SplitFactor[symbol]
            if symbol not in self.previousClose:                                        # Make a history call for symbol to get last closing price
                history = self.History(symbol, 1, Resolution.Daily) #, DataNormalizationMode.SplitAdjusted)
                if not history.empty:
                    history = history.close.unstack(0)[symbol]
                    if not history.empty:
                        self.previousClose[symbol] = history[0]
        for security in changes.RemovedSecurities:                                      # Remove symbols from previous close as they are removed from the universe
            symbol = security.Symbol
            self.previousClose.pop(symbol, None)
                
    def Rebalance(self):
        percentChange = {}                                                              # Dictionary to keep track of percent change from last close
        priceoverTwo = {}
        for symbol, previousClose in self.previousClose.items():                        # Populate Dictionary
            ## if self.Splits.ContainsKey(symbol):
            ##     continue
            if self.CurrentSlice.ContainsKey(symbol):
                ## self.Spliteventafter[symbol] = self.Value.SplitFactor[symbol]
                ## if self.Spliteventbefore[symbol] == self.Spliteventafter[symbol]:
                price = self.CurrentSlice[symbol].Close
                change = price/previousClose
                percentChange[symbol] = change
                priceoverTwo[symbol] = price
            symbols = list(percentChange.keys())                                        # Symbols under consideration
            sortedSymbols = sorted([x for x in symbols if percentChange[x] > 1 and priceoverTwo[x] > 1], key=lambda x : percentChange[x], reverse = True) # True is Highest first
            selected = sortedSymbols[:5]                                                # Get top xx symbols
        for symbol in selected:
            price = self.Securities[symbol].Price
            self.MarketOrder(symbol, self.cashused/price)                              #self.StopMarketOrder(symbol, -self.cashused/price, price*1.2) # Stop loss 20% higher than purchase price
    
    def LiquidatePositions(self):
        self.Liquidate()                                                                # Liquidate portfolio
    def OnMarketClose(self):
        for symbol in self.previousClose:                                               # Store new previous close values
            if self.CurrentSlice.ContainsKey(symbol):
                self.previousClose[symbol] = self.CurrentSlice[symbol].Close