clr.AddReference('QuantConnect.Research')
from QuantConnect.Research import QuantBook


class TachyonMultidimensionalChamber(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 11, 15)  # Set Start Date
        self.SetEndDate(2020, 12, 10)  # Set End Date
        self.SetCash(400000)  # Set Strategy Cash
        #self.AddUniverse(self.CoarseSelectionFunction)
        
        self.UniverseSettings.ExtendedMarketHours = True
        self.UniverseSettings.Resolution = Resolution.Minute
        symbols = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA) ]
        self.AddUniverseSelection( ManualUniverseSelectionModel(symbols) )
        
        #self.SetSecurityInitializer(self.SecurityInitializer)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        
        self.sd = {}
    
    
    def OnSecuritiesChanged(self, changed):
        for security in changed.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.sd:
                self.sd[symbol] = SymbolData(self, symbol)
        for security in changed.RemovedSecurities:
            symbol = security.Symbol
            sd = self.sd.pop(symbol, None)
            if sd:
                sd.dispose()
    
    
    #reset VWAP 
    def OnEndOfDay(self):
        for s in self.sd:
            self.sd[s].vwap.Reset()
    
        self.Log("OnEndOfDay")
    
    
    ########### 
    # on data #
    ###########
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''

        if self.IsWarmingUp:
            return

        if data.Time.minute == 0:
            for symbol, symbol_data in self.sd.items():
                self.Plot("VWAP", str(symbol), symbol_data.vwap.Current.Value)
       

        
class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.vwap = VolumeWeightedAveragePriceIndicator(20) #60*24 = 1440 minutes in a day
        self.symbol = symbol
        
        hist = algorithm.History(symbol, 10, Resolution.Minute).loc[symbol]
        for idx, bar in hist.iterrows():
            tradeBar = TradeBar(idx, symbol, bar.open, bar.high, bar.low, bar.close, bar.volume, timedelta(minutes=1))
            self.vwap.Update(tradeBar)
    
        self.consolidator = TradeBarConsolidator(Resolution.Minute)
        self.consolidator.DataConsolidated += self.consolidation_handler
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)

    def consolidation_handler(self, sender, consolidated):
        self.vwap.Update(consolidated)
        
    def dispose(self):
        algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator)