import talib
import numpy as np

class VerticalNadionComputer(QCAlgorithm):
    
    period = 21
    leverage = 1
    assets = []
    symbol_data_by_symbol = {}

    def Initialize(self):
        self.SetStartDate(2020, 5, 22)
        self.SetCash(100000)
        
        for ticker in ['SPY', 'QQQ', 'TLT']:
            symbol = self.AddEquity(ticker, Resolution.Minute).Symbol
            self.assets.append(symbol)
            self.symbol_data_by_symbol[symbol] = SymbolData(self, symbol, self.period)
        
        self.bond = self.AddEquity('IEF', Resolution.Minute).Symbol
        
        self.Schedule.On(self.DateRules.EveryDay(self.bond), self.TimeRules.AfterMarketOpen(self.bond, 65), self.trade)


    def trade(self):
        wt_a = 0; wt = self.leverage / len(self.Securities)
        
        for symbol in self.assets:
            symbol_data = self.symbol_data_by_symbol[symbol]
            AroonDown, AroonUp = talib.AROON(symbol_data.highs, symbol_data.lows, self.period)
            
            if AroonUp[-1] > AroonDown[-1]:
                self.SetHoldings(symbol, wt)
                wt_a += wt
            else:
                self.Liquidate(symbol)
        
        wt_b = self.leverage - wt_a
        
        self.SetHoldings(self.bond, wt_b)
        
            
class SymbolData:
    highs = np.array([])
    lows = np.array([])
    
    def __init__(self, algorithm, symbol, period):
        self.symbol = symbol
        self.algorithm = algorithm
        self.period = period
        
        # setup consolidator
        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.consolidation_handler
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
        history = algorithm.History(symbol, period + 1, Resolution.Daily)
        if history.empty or 'high' not in history.columns:
            return
        for time, row in history.loc[symbol].iterrows():
            self.highs = np.append(self.highs, row.high)
            self.lows = np.append(self.lows, row.low)
            
    
    @property
    def IsReady(self):
        return len(self.highs) == self.period + 1 and len(self.lows) == self.period + 1
        
    def consolidation_handler(self, sender, consolidated):
        self.highs = np.append(self.highs, consolidated.High)[-(self.period+1):]
        self.lows = np.append(self.lows, consolidated.Low)[-(self.period+1):]