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, 11, 17)  # 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
        
        for s in self.Securities:
            self.Debug(self.Securities[s])
        
        #variables to keep track of
        self.sd = {} #all symbol data
        
        self.hod = {}
        self.lod = {}
        self.open = {}

        self.buffer = 0 #only enter trades after N buffer after open
        self.ema5_past = {}
        self.oh_past5 = {}
        
        #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_past5 is index 0, furthest is [-1]
    def getStop(self, oh_past5, 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_past5[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_past5[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()
    
    
    ########### 
    # 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

        tradeBars = data.Bars #OHLC of past time interval
    
        for s in self.sd:
            
            self.Debug(self.Time.weekday())
            #reset VWAP on new day
            if self.day != self.Time.weekday(): #new day
                self.day = self.Time.weekday()
                self.vwap = self.VWAP(s, 2000, Resolution.Minute)
            
            #highs and lows of candles in past 5 minutes (even premarket)
            if data.ContainsKey(s):
                if s not in self.oh_past5:
                    self.oh_past5[s] = [] 
                if len(self.oh_past5[s]) < 5:
                    self.oh_past5[s].append(str(tradeBars[s].High) + " " + str(tradeBars[s].Low))
                elif len(self.oh_past5[s]) == 5:
                    del self.oh_past5[s][-1]
                    self.oh_past5[s].insert(0, str(tradeBars[s].High) + " " + str(tradeBars[s].Low))
        
            if self.IsMarketOpen(s) and data.ContainsKey(s):
                #this block should not execute in practice since PM data should always be used
                if s not in self.hod:
                    self.hod[s] = -1.0
                    self.lod[s] = 2000000000.0
                    self.open[s] = -1.0
                
                
                #get hod and lod
                if tradeBars[s].High > self.hod[s]:
                    self.hod[s] = tradeBars[s].High
                if tradeBars[s].Low < self.lod[s]:
                    self.lod[s] = tradeBars[s].Low
                if self.open[s] == -1.0:
                    self.open[s] = tradeBars[s].Open
                
                
                ################
                # 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)
                hod = float(self.hod[s])
                lod = float(self.lod[s])
                ema5_past = self.ema5_past[s]
                day_range = hod - lod
                
                #position size in dollars
                size = 50.0
                
                if self.Portfolio[s].Invested == False and self.Portfolio.MarginRemaining > size and self.buffer >= 1:
                    if ema5 > vwap and (price > (lod + (.45 * day_range))) and ema5 > ema5_past and (price - vwap) < atr5 and self.Time.hour < 15: #and day_range <  day_atr and price > vwap:
                        #long
                        self.Debug(self.Time)
                        self.Debug("New long")
                        
                        stop = self.getStop(self.oh_past5[s], vwap, atr5)
                        shares = self.positionSize(stop, price, size)
                        
                        self.Debug("margin remaining: " + str(self.Portfolio.MarginRemaining))
                        self.Debug(s)
                        #self.Debug("nShares: " + str(shares))
                        self.Debug("currPrice: " + str(price))
                        self.Debug("stop: " + str(stop))
                        self.Debug("buying power used: " + str(shares*price))
                        self.Debug("VWAP: " +  str(vwap))
                        self.Debug("")
                        
                        self.MarketOrder(s, shares)
                        stop = -1*shares
                        self.StopMarketOrder(s, stop, lod)
                elif self.Portfolio[s].Invested == True:
                    #a "good" trade
                    if ema5 < ema9 and price < ema9 and ema5 < ema5_past:
                        self.Liquidate(s)
                        
                        self.Debug(self.Time)
                        self.Debug("Sold " + str(s))
                    #is failing, and stop out before hitting S/L
                    
                    #has consolidated for too long without moving up
                    
                    #take partials:
                        #HOD (or other significant S/R) and PM high at least are reasonable
                        #>daily ATR
                        
                
                #close all at 3:55 if any positions are open
                #should change this to time to market close
                    #use BeforeMarketClose()
                if self.Time.hour >= 15 and self.Time.minute >= 55: 
                    self.Liquidate()
                
                #increase buffer during market hours
                self.buffer += 1
                
            #not during market hours, reset various
            else:
                self.hod[s] = -1.0
                self.lod[s] = 2000000000.0
                self.open[s] = -1.0
                self.buffer = 0
            
            
            #past period ema5
            self.ema5_past[s] = self.sd[s].ema5.Current.Value
       

        
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)