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

class TachyonMultidimensionalChamber(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 12, 5)  # Set Start Date
        self.SetEndDate(2020, 12, 18)  # Set End Date
        self.SetCash(400000)  # Set Strategy Cash
        self.AddUniverse(self.CoarseSelectionFunction)
        
        self.SetSecurityInitializer(self.SecurityInitializer)
        
        self.UniverseSettings.ExtendedMarketHours = True
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.UniverseSettings.Leverage = 4
        #self.UniverseSettings.Resolution = Resolution.Hour #can comment/change this out
        self.SetWarmUp(5)
        
        #for s in self.Securities:
        #    self.Debug(self.Securities[s])
        
        
        
        ###variables to keep track of
        self.sd = {} #all symbol data
        
        self.pastNoPre = {}
        self.pastWithPre = {}
        
        self.keepPast = 500 #how many past periods (1 minute) to keep
        
        self.hod = {}
        self.lod = {}
        self.open = {}

        self.buffer = 0 #only enter trades after N buffer after open
        self.ema5_past = {}
        
        #self.pmhigh = {}
        #self.pmlow = {}
        
 
        
        #self.day = self.Debug(self.Time.weekday())

        


    def SecurityInitializer(self, security):
        security.SetLeverage(4)
    
    
    def CoarseSelectionFunction(self, universe):  
        selected = []
        for coarse in universe:  
            if coarse.Volume > 70000000 and coarse.Value > 10 and coarse.HasFundamentalData:
                symbol = coarse.Symbol
                selected.append(symbol)
        return selected #list of objects of type Symbol
    
    
    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
            self.sd.pop(symbol, None)



    ################
    # my functions #
    ################
    ###sizing
    def positionSize(self, stop, currPrice, dollarSize):
        nShares = int(dollarSize / abs(stop - currPrice))
        return nShares
    
        
    ###get stop
    #most recent of oh_past is index 0, furthest is [-1]
    def getStop(self, oh_past, vwap, atr):
        #init and parse
        firstToCross = -1000
        low_postCross = 2000000000.0
        
        #get first candle that crossed
        for i in range(0, 5):
            high, low = oh_past[i].split()
            high = float(high)
            if high > vwap:
                firstToCross = i
        
        #low of all candles post-cross
        for i in range(0, i):
            high, low = oh_past[i].split()
            low = float(low)
            if low < low_postCross:
                low_postCross = low
    
        return low_postCross + (.5*atr)
    
    
    ###reset VWAP 
    def OnEndOfDay(self):
        for s in self.sd:
            self.sd[s].vwap.Reset()
    
    
    ###consolidate bars
    def consolidate(oneminutes, period, self):
        
        consolidated = []
        curr = 0
        
        currOpen = 0
        currHigh = 0
        currLow = 2000000000
        currClose = 0
        currVol = 0
        
        self.Debug("oneminutes")
        self.Debug(oneminutes)
        
        self.Debug("period")
        self.Debug(period)
        
        self.Debug("self")
        self.Debug(self)
        
        for i in range(0, len(oneminutes)):
            curr += 1
            o, h, l, c, v = oneminutes.split() #ignition bar
            o = float(o)
            h = float(h)
            l = float(l)
            c = float(c)
            v = float(v)
            
            #init
            if curr == 1:
                currOpen = o
            
            #update
            currVol += v
            
            if h > currHigh:
                currHigh = h
            if low < currLow:
                currLow = l
            
            #have filled period up
            if curr == period:
                currClose = c
                toappend = str(currOpen) + " " + str(currHigh) + " " + str(currLow) + " " + str(currClose) + " " + str(currVol)
                
                consolidated.append(toappend)
                
                #reset
                currOpen = 0
                currHigh = 0
                currLow = 2000000000
                currClose = 0
                currVol = 0
                
                curr = 0
            
        return consolidated
        
        
    
    ########### 
    # 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
        '''
        

        tradeBars = data.Bars #OHLC of past time interval
    
    
        for s in self.sd:
            curr_ohlcv = str(tradeBars[s].Open) + " " + str(tradeBars[s].High) + " " + str(tradeBars[s].Low) + " " + str(tradeBars[s].Close) + " " + str(tradeBars[s].Volume)
            
            
            #calc and add range of past 500 candles, with premarket
            if data.ContainsKey(s):            
                currRange = tradeBars[s].High - tradeBars[s].Low
                if s not in self.pastWithPre:
                    self.pastWithPre[s] = []
                    self.pastWithPre[s].append(curr_ohlcv)
                else:
                    if len(self.pastWithPre[s]) >= 500:
                        self.pastWithPre[s].pop()
                    self.pastWithPre[s].insert(0, curr_ohlcv)
                
                #warming up
                if self.IsWarmingUp:
                    return
            
            
            #market is open now
            if data.ContainsKey(s) and self.IsMarketOpen(s):
                
                #save past 500
                if s not in self.pastNoPre:
                    self.pastNoPre[s] = []
                    self.pastNoPre[s].append(curr_ohlcv)
                else:
                    if len(self.pastNoPre[s]) >= 500:
                        self.pastNoPre[s].pop()
                    self.pastNoPre[s].insert(0, curr_ohlcv)
            
        
                ################
                # trade setups #
                ################
                ### consolidation above VWAP, short-term trend up, has not traveled much of a daily ATR, in upper part of daily range --> ORB up
                vwap = self.sd[s].vwap.Current.Value
                ema5 = self.sd[s].ema5.Current.Value
                ema9 = self.sd[s].ema9.Current.Value
                atr5 = self.sd[s].atr5.Current.Value
                #day_atr = self.atrDaily[d].Current.Value
                
                price = float(tradeBars[s].Close)
                
                #position size in dollars
                size = 50.0
                
                
                if s in self.pastWithPre and len(self.pastWithPre[s]) == self.keepPast: #ready
                    if self.Time.hour >= 15 and self.Time.minute >= 55:
                            self.Debug(s)


#                            #consolidate
#                            consolidated = self.consolidate(self.pastWithPre[s], self)
                        
            

        
class SymbolData:
    def __init__(self, algorithm, symbol):
        #algorithm.VWAP(symbol, 10000, Resolution.Minute)
        self.vwap = algorithm.VWAP(symbol, 2000, Resolution.Minute) #60*24 = 1440 minutes in a day
        self.ema5 = algorithm.EMA(symbol, 5, Resolution.Minute)
        self.ema9 = algorithm.EMA(symbol, 9, Resolution.Minute)
        self.atr5 = algorithm.ATR(symbol, 5, Resolution.Minute)
        
        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.ema5.Update(idx, bar.close)
            self.ema9.Update(idx, bar.close)
            self.atr5.Update(tradeBar)
            self.vwap.Update(tradeBar)