from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
class TransdimensionalParticleThrustAssembly(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 7, 1)                                                  # Set Start Date
        self.SetEndDate(2020, 7, 7)                                                    # 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.SetSecurityInitializer(self.CustomSecurityInitializer)
        self.UniverseSettings.FillForward = False
        self.UniverseSettings.ExtendedMarketHours = True
        # self.UniverseSettings.Leverage = 1.0
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.2)) 
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(9, 0), self.Rebalance)                  # Scheduled Events
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(9, 1), self.Rebalance_Second)
        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.previous_d_close = {}                                                          # Dictionary to keep track of previous close for each symbol
        self.filtered = []
        self.donottrade = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in []] #, '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 5 > 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 < 5000000000]
        
    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
            if symbol not in self.previous_d_close:                                         # 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.previous_d_close[symbol] = history[0]
        for security in changes.RemovedSecurities:                                          # Remove symbols from previous close as they are removed from the universe
            self.previous_d_close.pop(security.Symbol, None)
                
    def Rebalance(self):
        percent_change = {}                                                                 # Dictionary to keep track of percent change from last close
        price_restriction = {}
        for symbol, previous_d_close in self.previous_d_close.items():                      # Populate Dictionary
            if self.CurrentSlice.Bars.ContainsKey(symbol):
                last_price = self.CurrentSlice.Bars[symbol].Close
                change = last_price/previous_d_close
                percent_change[symbol] = change
                price_restriction[symbol] = last_price
            symbols = list(percent_change.keys())                                           # Symbols under consideration
            sorted_symbols = sorted([x for x in symbols if percent_change[x] > 1.0 and price_restriction[x] > 2], key=lambda x : percent_change[x], reverse = True) # True is Highest first
            self.filtered = sorted_symbols[:10]                                             # Get top xx symbols
    
    def Rebalance_Second(self):
        price_above_ma = {}
        price_below_max = {}
        percent_change = {}                                                                 # Dictionary to keep track of percent change from last close
        selected_symbols = []
        for symbol, previous_d_close in self.previous_d_close.items():
            if symbol in self.filtered:
                if self.CurrentSlice.Bars.ContainsKey(symbol):
                    history_data_max = self.History(symbol, 60, Resolution.Daily).high.unstack(level=0)[symbol]
                    history_data_close = self.History(symbol, 60, Resolution.Daily).close.unstack(level=0)[symbol]
                    forty_d_max = history_data_max[-60:-2].max()
                    forty_d_avg = history_data_close[-60:-2].mean()
                    price_above_ma[symbol] = forty_d_max / forty_d_avg
                    current_price = self.CurrentSlice.Bars[symbol].Close
                    price_below_max[symbol] = current_price / forty_d_max
                    percent_change[symbol] = current_price/previous_d_close
                    symbols = list(price_above_ma.keys())
                    selected_symbols = sorted([x for x in symbols if price_above_ma[x] > 0.8 and price_below_max[x] < 5], key=lambda x : percent_change[x], reverse = True)#[:1]
        for symbol in selected_symbols:
            price = self.Securities[symbol].Price
            self.LimitOrder(symbol, -self.cashused/price, (price * 0.99))
            #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 CustomSecurityInitializer(self, security):
        security.SetDataNormalizationMode(DataNormalizationMode.SplitAdjusted)
        security.SetFeeModel(CustomFeeModel())
        #security.SetSlippageModel(CustomSlippageModel(self))
        
    def OnMarketClose(self):
        for symbol in self.previous_d_close:                                                # Store new previous close values
            if self.CurrentSlice.ContainsKey(symbol):
                self.previous_d_close[symbol] = self.CurrentSlice[symbol].Close
    
class CustomFeeModel():                                                                     # Slippage and Fees together
    def GetOrderFee(self, parameters):
        loss_total = (parameters.Security.Price * 0.011 * parameters.Order.AbsoluteQuantity) + (parameters.Order.AbsoluteQuantity * 0.005)
        fee = max(1, loss_total)
        return OrderFee(CashAmount(fee, 'USD'))

#class CustomSlippageModel:
#    def __init__(self, algorithm):
#        self.algorithm = algorithm
#    def GetSlippageApproximation(self, asset, order):                                      # custom slippage math
#        slippage = np.float(asset.Price) * 0.02 #asset.Price * d.Decimal(0.0001 * np.log10(2*float(order.AbsoluteQuantity)))
#        return slippage