# VXX version - best length 22 
# VIX hour version - best length 11
# VIX daily version - best length 22

import numpy as np
from datetime import datetime


class BasicTemplateAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2021, 7, 17)  
        #self.SetEndDate(2019, 3, 17) 
        self.SetEndDate(datetime.now())    
        self.SetCash(25000)           
        self.Settings.FreePortfolioValuePercentage = 0.00
        self.data = {}
        self.SetBenchmark("SPY")
        #period = 10*21
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.trailing_stop = self.GetParameter("trailing-stop")
        self.trailing_stop = float(self.trailing_stop) if self.trailing_stop else 0.04
        
        self.moving_average = self.GetParameter("moving-average")
        self.moving_average = float(self.moving_average) if self.moving_average else 13
        
        self.rsi_value = self.GetParameter("rsi-value")
        self.rsi_value = float(self.rsi_value) if self.rsi_value else 51
        
        self.sma_tolerance = self.GetParameter("sma-tolerance")
        self.sma_tolerance = float(self.sma_tolerance) if self.sma_tolerance else 0.0063
        
        self.vix_length = self.GetParameter("vix-length")
        self.vix_length = float(self.vix_length) if self.vix_length else 8
        
        self.rsi_upper = self.GetParameter("rsi-upper")
        self.rsi_upper = float(self.rsi_upper) if self.rsi_upper else 88
        
        
        self.screener_price = self.GetParameter("screener-price")
        self.screener_price = float(self.screener_price) if self.screener_price else 60
        
        self.AddRiskManagement(TrailingStopRiskManagementModel(self.trailing_stop))

        self.SPY = self.AddEquity("SPY", Resolution.Hour).Symbol
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose(self.SPY,20),       
                 self.StopTrading)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen(self.SPY, 0),       
                 self.StartTrading)
        
        self.SPY = self.AddEquity("SPY", Resolution.Hour).Symbol

        
        self.vix = self.AddEquity("VXX", Resolution.Minute).Symbol
        self.staticAssets = [self.SPY, self.vix]
        self.vixSma = self.SMA(self.vix, self.vix_length, Resolution.Hour)
        vixHistory = self.History(self.vix, 11, Resolution.Hour)
        for tuple in vixHistory.loc[self.vix].itertuples():
            self.vixSma.Update(tuple.Index, tuple.close)
            
        self.Log(". VIX SMA INITIALIZED: " + str(self.vixSma.Current.Value)) # + ". VIX INITIALIZED: " + str(self.vix.Current.Value))
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.BeforeMarketClose(self.SPY,20),       
                 self.StopTrading)
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen(self.SPY, 0),       
                 self.StartTrading)
                 
        # self.Schedule.On(self.DateRules.EveryDay("SPY"),
        #          self.TimeRules.AfterMarketOpen(self.SPY, 30),       
        #          self.SafetySwitch)
    
        self.AddUniverse(self.CoarseFilter, self.FineFilter)
        
        '''
        If VVIX > 50 day SMA, 
            Sell all current assets
            Switch to vxx and IEF (50/50)
        '''
        
        self.UniverseSettings.Resolution = Resolution.Hour
        self.lastMonth = -1
        self.lastHour = -1
        self.allowTrades = True
        self.switchSafety = False
        
        self.Log("Initialized")
        
    def StartTrading(self):
        self.switchSafety = False
        self.allowTrades = True
        
        # if self.Securities[self.vix].Price > self.vixSma.Current.Value:
        #     self.switchSafety = True
        #     self.allowTrades = False
        #     self.Log("StartTrading: allowtrades = False, switchSafety = True")
        # else:
        #     self.switchSafety = False
        #     self.allowTrades = True
        #     self.Log("StartTrading: allowtrades = True, switchSafety = False")
            
    # def SafetySwitch(self):
    #     if self.switchSafety == False:
    #         self.Liquidate(self.vxx)
    #         self.Liquidate(self.ief)
    #         self.Log("switchSafety is False, liquidating vxx and ief")
    #     elif self.switchSafety == True:
    #         self.Log("switchSafety is True now liquidating")
    #         if not self.Portfolio[self.vxx].Invested and not self.Portfolio[self.ief].Invested:
    #             self.Liquidate()
    #         #self.SetHoldings(self.vxx, 0.1)
    #         #lf.SetHoldings(self.ief, 0.5)
    
    
    # def SafetySwitch(self):
    #     if 1==1:
    #         self.Liquidate()
        
    def StopTrading(self):
        self.Log("Stopping Trading")
        self.allowTrades = False
        
    def CoarseFilter(self, coarse):
        if self.lastMonth == self.Time.month:
            return Universe.Unchanged
        self.lastMonth = self.Time.month
        
        topStocksByVolume = sorted([x for x in coarse if x.Price > self.screener_price and x.Volume > 0], key = lambda x: x.DollarVolume, reverse=True)[:6]
        return [x.Symbol for x in topStocksByVolume]
    
    def FineFilter(self, fine):
        
        topFiveMarketCap = sorted([x for x in fine], key = lambda x: x.MarketCap, reverse = True)[:5]
        
        ''' Manually exclude companies '''
        filteredSymbols = [x.Symbol for x in topFiveMarketCap if x.Symbol.Value not in ["GME", "AMC", "GOOG"]]
        
        ''' Filter out high volatile stocks '''
        preSdSymbols = [x for x in filteredSymbols]
        history = self.History(preSdSymbols, 50, Resolution.Daily)
        standardDeviations = {}
        for symbol in preSdSymbols:
            sd = StandardDeviation(50)
            for tuple in history.loc[symbol].itertuples():
                sd.Update(tuple.Index, tuple.close)
            standardDeviations[symbol] = sd.Current.Value
        finalSymbols = [x for x in preSdSymbols if standardDeviations[x] < 1000]#5]
        
        
        self.Log("Picked number of symbols in universe: " + str(len(finalSymbols)))
        
        return finalSymbols
        
    def PlaceTrades(self, data):
        isUptrend = []
        for symbol, symbolData in self.data.items():
            self.Debug(str(symbol.Value) + " at " + str(self.Time) + ". RSI: " + str(symbolData.Rsi.Current.Value) + ". SMA: " + str(symbolData.Sma.Current.Value) + ". PRICE: " + str(self.Securities[symbol].Price))
            if not data.ContainsKey(symbol):
                self.Log("Does not contain data for " + str(symbol))
                continue
            if self.Securities[symbol].Price > (symbolData.Sma.Current.Value * (1 + self.sma_tolerance)) and not symbolData.Rsi.Current.Value < self.rsi_value and not symbolData.Rsi.Current.Value > self.rsi_upper and not self.Securities[self.vix].Price > self.vixSma.Current.Value:
                isUptrend.append(symbol)
            # elif self.Portfolio[symbol].Invested:
            #     self.Log("Liquidating: " + str(symbol))
            #     self.Liquidate(symbol, "SMA: " + str(self.data[symbol].Sma.Current.Value) + ". RSI: " + str(self.data[symbol].Rsi.Current.Value))
        
        for symbol in isUptrend:
            self.Debug("Buying: " + str(symbol))
            self.SetHoldings(symbol, 1/len(isUptrend), False, "SMA: " + str(self.data[symbol].Sma.Current.Value) + ". RSI: " + str(self.data[symbol].Rsi.Current.Value))
        
    def OnData(self, data):
        if self.IsWarmingUp: return
        self.Log(". VIX: " + str(self.Securities[self.vix].Price) + ". VIX SMA: " + str(self.vixSma.Current.Value))
        # self.Log(". Symbols: " + str(self.data.keys))
        if self.allowTrades == False: return
        #i moved place trades liquidate logic back to ondata
        if self.Securities[self.vix].Price > self.vixSma.Current.Value:
            self.Liquidate()
            self.Log("VIX CHECK FAILED LIQUIDATING")
        if self.lastHour == self.Time.hour or self.Time.hour < 10:
            return   
        self.lastHour = self.Time.hour
        self.PlaceTrades(data)
            
    def OnSecuritiesChanged(self, changes):
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            if symbol in self.staticAssets: continue
            added.MarginModel = PatternDayTradingMarginModel()
            sma = self.SMA(symbol, self.moving_average, Resolution.Hour)
            rsi = self.RSI(symbol, 14, MovingAverageType.Simple, Resolution.Hour)
            # sto = self.STO(symbol)
            history = self.History(symbol, 15, Resolution.Hour)
            for tuple in history.loc[symbol].itertuples():
                sma.Update(tuple.Index, tuple.close)
                rsi.Update(tuple.Index, tuple.close)
                # sto.Update(TradeBar)
            
            self.Log(f'New Securities Added: {[security.Symbol.Value for security in changes.AddedSecurities]}')
            
            
            symbolData = SymbolData(sma, rsi)
            self.data[symbol] = symbolData
            
            #self.data[symbol].Sma.Current.Value
            
        for removed in changes.RemovedSecurities:
            symbol = removed.Symbol
            self.Liquidate(symbol)
            self.data.pop(symbol, None)
            self.Log(f'Securities Removed{[security.Symbol.Value for security in changes.RemovedSecurities]}')
            
class SymbolData:
    def __init__(self, sma, rsi):
        self.Sma = sma
        self.Rsi = rsi
        # self.Sto = sto
            
            
# FOR JOVAD: Need 2 versions 1 with a universe that only trades top 5 market cap stocks and 1 with Dropbox spreadsheet
# Need to add STO, RSI, AND HAVE 3 SMA HERE (fast, medium, slow)
# Need minute support with bars, 60 minutes consolidated for now
# Need to understand how I can add more indicators if I want - struggled with warming up onData. Tried adding more indicators using for symbol, sma in self.data.items(): but it didn't work

# FOR JOVAD NEW SESSION: HOW TO PICK STOCKS WITH HIGHEST 


# FOR JOVAD 6/29- NEW STRATEGY TO TEST: IF RSI PREVIOUS CLOSE < RSI CURRENT CLOSE ENTER TRADE ELSE EXIT TRADE, TRAILING STOP 5% HOURLY CHART
# STOCK PICKING HOUR SMA NEEDS TO BE ADJUSTED TO CRYPTO AS WELL WITHOUT UNIVERSE SELECTION + FIX EXCHANGE IS CLOSED ERROR
# ALSO FOR STOCK PICKING - WE NEED TO ADD FUNCTIONALITY TO REMOVE MANUALLY SELECTED TICKERS LIKE GOOG AND GME
# ASK FOR ADVISE IF WE CAN FIND A LOGIC TO EXCLUDE DANGEROUS TICKERS LIKE AMC AND GME