# 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, 10, 15)   
        #self.SetEndDate(2021, 10, 5) 
        self.SetEndDate(datetime.now())    
        self.SetCash(23400)           
        self.Settings.FreePortfolioValuePercentage = 0.00
        self.data = {}
        self.SetBenchmark("SPY")
        
        self.CoarseSize = 8
        #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 55
        
        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 65
        
        self.qqq_length = self.GetParameter("qqq-length")
        self.qqq_length = float(self.qqq_length) if self.qqq_length else 20
        
        
        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,2),       
        #          self.ForceTrade)
                 
        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, 2),       
        #          self.ForceTrade)
                 
                 
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose(self.SPY, 1), self.EndOfDay)
        

        self.vix = self.AddEquity("VXX", Resolution.Minute).Symbol
        self.qqq = self.AddEquity("QQQ", Resolution.Minute).Symbol
        self.staticAssets = [self.SPY, self.vix, self.qqq]
        self.vixSma = self.SMA(self.vix, self.vix_length, Resolution.Hour)
        self.vixRsi = self.RSI(self.vix, 14, Resolution.Hour)
        self.qqqSma = self.SMA(self.qqq, self.qqq_length, Resolution.Hour)
        
        vixHistory = self.History(self.vix, 14, Resolution.Hour)
        for tuple in vixHistory.loc[self.vix].itertuples():
            self.vixSma.Update(tuple.Index, tuple.close)
            self.vixRsi.Update(tuple.Index, tuple.close)
            

        qqqHistory = self.History(self.qqq, 20, Resolution.Hour)
        for tuple in qqqHistory.loc[self.qqq].itertuples():
            self.qqqSma.Update(tuple.Index, tuple.close)
            

            
        self.Log(". VIX SMA INITIALIZED: " + str(self.vixSma.Current.Value)) 
        
        self.AddUniverse(self.CoarseFilter)
        '''
        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 EndOfDay(self):
        if self.IsWarmingUp: return
        self.Log('Checking End Of Day VIX & QQQ')
        if self.Securities[self.qqq].Price < self.qqqSma.Current.Value:
            #self.Liquidate()
            self.Log('Daily VIX & QQQ Check Failed')
            self.Log("Stopping Trading")
            self.allowTrades = False
        else:
            self.Log('Daily VIX & QQQ Check Passed')
            self.Log("Stopping Trading")
            self.allowTrades = False
        
    def StartTrading(self):
        self.switchSafety = False
        self.allowTrades = True
    
    # def ForceTrade(self):
    #     if self.IsWarmingUp: return
    #     self.Log("FINAL TRADE")
    #     # 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 and self.vixRsi.Current.Value > 50 or self.Securities[self.qqq].Price < self.qqqSma.Current.Value:
    #         self.Liquidate()
    #         self.Log("Force Trade VIX or QQQ Check Failed, Liquidating")
    #     if self.lastHour == self.Time.hour or self.Time.hour < 10:
    #         return   
    #     self.lastHour = self.Time.hour
    #     self.PlaceTrades(data)
    
    def CoarseFilter(self, coarse):
        #if self.lastMonth == self.Time.month:
        self.Log(f'CoarseFilter called on {self.Time}')
            #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 and x.HasFundamentalData and
                x.Symbol.Value not in ["GME", "AMC", "GOOG", "SPY"]], # Manually exclude companies
            key = lambda x: x.DollarVolume, reverse=True)[:self.CoarseSize]
            
        return [x.Symbol for x in topStocksByVolume]
    
    # def FineFilter(self, fine):
        
    #     finalSymbols = sorted([x for x in fine], key = lambda x: x.MarketCap, reverse = True)[:8]
        
    #     ''' Filter out high volatile stocks '''
    #     # preSdSymbols = [x for x in coarse]
    #     # 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 [x.Symbol for x in 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*T: " + str(symbolData.Sma.Current.Value * (1 + self.sma_tolerance)) + ". 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.qqq].Price < self.qqqSma.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 and self.vixRsi.Current.Value > 50:
            self.Liquidate()
            self.Log("ONDATA VIX OR QQQ 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, tag="Symbol Removed From Universe")
            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