# 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.Minute).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]
        
        
        # 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.Minute
        self.lastMonth = -1
        self.lastHour = -1
        self.allowTrades = True
        self.switchSafety = False
        
        self.Log("Initialized")
        
        self.isUptrend = []
        self.upTrendCount = 0
        
    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 OnData(self, data):
        if self.upTrendCount > 0:
            self.Debug(f"Placing trades at {self.Time}")
            self.size = len(self.isUptrend)
            for symbol in self.isUptrend:
                self.Debug("Buying: " + str(symbol))
                self.SetHoldings(symbol, 1/len(self.isUptrend), False)#, "SMA: " + str(self.data[symbol].Sma.Current.Value) + ". RSI: " + str(self.data[symbol].Rsi.Current.Value))
            self.upTrendCount = 0
            self.isUptrend = []
            
    def OnSecuritiesChanged(self, changes):
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            if symbol in self.staticAssets: 
                symbolData = SymbolData(symbol, self)
                self.data[symbol] = symbolData
                continue
            added.MarginModel = PatternDayTradingMarginModel()
            
            self.Log(f'New Securities Added: {[security.Symbol.Value for security in changes.AddedSecurities]}')
            
            symbolData = SymbolData(symbol, self)
            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, symbol, algorithm):
        self.Symbol = symbol
        self.algorithm = algorithm
        
        hourlyConsolidator = TradeBarConsolidator(self.Custom)
        hourlyConsolidator.DataConsolidated += self.OnDataConsolidated
        algorithm.SubscriptionManager.AddConsolidator(symbol, hourlyConsolidator)
        
        # Define our indicator
        self.Sma = SimpleMovingAverage(int(self.algorithm.moving_average))
        self.Rsi = RelativeStrengthIndex(14, MovingAverageType.Simple)
        # Register indicator to our consolidator
        algorithm.RegisterIndicator(symbol, self.Sma, hourlyConsolidator)
        algorithm.RegisterIndicator(symbol, self.Rsi, hourlyConsolidator)
        
        # Rolling window to hold 30 min bars
        self.barWindow = RollingWindow[TradeBar](2)
        
        
    # Store every thirty minute bar in rolling window
    def OnDataConsolidated(self, sender, bar):
        self.barWindow.Add(bar)

        if bar.Symbol in self.algorithm.staticAssets or not self.IsReady: return

        self.algorithm.Log(". VIX: " + str(self.algorithm.Securities[self.algorithm.vix].Price) + ". VIX SMA: " + str(self.algorithm.data[self.algorithm.vix].Sma.Current.Value) + ". VIX RSI: " + str(self.algorithm.data[self.algorithm.vix].Rsi.Current.Value))
        self.algorithm.Log(". Symbols: " + str(self.algorithm.data.keys))
        if self.algorithm.allowTrades == False: return
        #i moved place trades liquidate logic back to ondata
        if self.algorithm.Securities[self.algorithm.vix].Price > self.algorithm.data[self.algorithm.vix].Sma.Current.Value and self.algorithm.data[self.algorithm.vix].Rsi.Current.Value > 50:
            self.algorithm.Liquidate()
            self.algorithm.Log("ONDATA VIX OR QQQ CHECK FAILED LIQUIDATING")
        # if self.algorithm.lastHour == self.algorithm.Time.hour or self.algorithm.Time.hour < 10:
        #     return   
        self.algorithm.lastHour = self.algorithm.Time.hour
        self.PlaceTrades(bar)
                
    @property
    def IsReady(self):
        return self.Sma.IsReady and self.Rsi.IsReady and self.barWindow.IsReady
        
    def Custom(self, dt):
        period = timedelta(hours=1)
        start = dt.replace(minute=30)
        if start > dt:
            start -= period
        return CalendarInfo(start, period)
        
    def PlaceTrades(self, data):
        self.algorithm.Debug(str(self.Symbol.Value) + " at " + str(self.algorithm.Time) + ". RSI: " + str(self.Rsi.Current.Value) + ". SMA*T: " + str(self.Sma.Current.Value * (1 + self.algorithm.sma_tolerance)) + ". PRICE: " + str(self.algorithm.Securities[self.Symbol].Price))
        if self.algorithm.Securities[self.Symbol].Price > (self.Sma.Current.Value * (1 + self.algorithm.sma_tolerance)) and not self.Rsi.Current.Value < self.algorithm.rsi_value and not self.Rsi.Current.Value > self.algorithm.rsi_upper and not self.algorithm.Securities[self.algorithm.qqq].Price < self.algorithm.data[self.algorithm.qqq].Sma.Current.Value:
            self.algorithm.upTrendCount += 1
            self.algorithm.isUptrend.append(self.Symbol)