# 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

        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.AfterMarketOpen(self.SPY, 0),       
                 self.StartTrading)
                 
        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.allowTrades = False
        self.Log('End Of Day Stop Trade')
        
    def StartTrading(self):
        self.switchSafety = False
        self.allowTrades = True

    
    def CoarseFilter(self, coarse):
        self.Log(f'CoarseFilter called on {self.Time}')
        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]
        
        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)
        
        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) + ". VIX RSI: " + str(self.vixRsi.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)
            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)
            
            self.Log(f'New Securities Added: {[security.Symbol.Value for security in changes.AddedSecurities]}')
            
            symbolData = SymbolData(sma, rsi)
            self.data[symbol] = symbolData
            
        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