from datetime import datetime, date, time
from datetime import timedelta
import pandas as pd

class SymbolData:
    def __init__(self,algo,symbol):
        self.algo = algo
        self.symbol = symbol
        self.openPrice = -1   
        self.lowestPrice = -1
        self.ticket = None
        
        self.daily = RollingWindow[TradeBar](2)
        self.window = RollingWindow[TradeBar](2)
        
        # First RSI indicator and consolidator
        self.rsi1 = RelativeStrengthIndex( 14, MovingAverageType.Wilders)
        self.cons1 = TradeBarConsolidator(timedelta(minutes=15))
        self.algo.SubscriptionManager.AddConsolidator(symbol,self.cons1)
        self.algo.RegisterIndicator(symbol,self.rsi1,self.cons1)
        
        # First RSI indicator and consolidator
        self.rsi2 = RelativeStrengthIndex(14, MovingAverageType.Wilders)
        self.cons2 = TradeBarConsolidator(timedelta(minutes=2))
        self.algo.SubscriptionManager.AddConsolidator(symbol,self.cons2)
        self.algo.RegisterIndicator(symbol,self.rsi2,self.cons2)
       
        # Add consolidator for obtaining previous days close
        self.cons3 = TradeBarConsolidator(timedelta(days=1))
        self.cons3.DataConsolidated += self.OnDailyData
        self.algo.SubscriptionManager.AddConsolidator(symbol,self.cons3) 
        
        # Add daily bar to daily rolling window
    def OnDailyData(self, sender, bar):
        self.daily.Add(bar) 
        
        # Check Indicators are ready before initializing
    def IsReady(self):
        return self.rsi1.IsReady and self.rsi2.IsReady
    def dispose(self):
        self.algo.SubscriptionManager.RemoveConsolidator(self.symbol, self.cons1)
        self.algo.SubscriptionManager.RemoveConsolidator(self.symbol, self.cons2)
        self.algo.SubscriptionManager.RemoveConsolidator(self.symbol, self.cons3)
        
class TradeStrategyTest(QCAlgorithm):
    
   def Initialize(self):
        self.SetStartDate(2021,3, 1)  #Set Start Date
        self.SetEndDate(2021,4,1)    #Set End Date
        self.SetCash(30000)           #Set Strategy Cash
        self.SetWarmUp(12600)
        self.current = None
        
        # Set TimeZone
        self.SetTimeZone("America/New_York")
        self.dataBySymbol = {}
        self.rebalanceTime = datetime.min
        self.activeStocks = set()

        self.AddUniverse(self.CoarseFilter, self.FineFilter)
        self.UniverseSettings.Resolution = Resolution.Minute
        
   def CoarseFilter(self, coarse):
        # Rebalancing monthly
        if self.Time <= self.rebalanceTime:
            return self.Universe.Unchanged
        self.rebalanceTime = self.Time + timedelta(30)

        sortedByDollarVolume = sorted(coarse, key=lambda x: x.Volume, reverse=True)
        return [x.Symbol for x in sortedByDollarVolume if x.Price > 10 and x.Price <100
                                                and x.HasFundamentalData][:5] 
        
        
   def FineFilter(self, fine):
        
        sortedByPE = sorted(fine, key=lambda x: x.MarketCap,reverse = True)

        

        return [x.Symbol for x in sortedByPE if x.MarketCap > 0][:3]
        
        
        
       
        
   def OnSecuritiesChanged(self,changes):
       for x in changes.AddedSecurities:
           if x.Symbol not in self.dataBySymbol:
               self.dataBySymbol[x.Symbol] = SymbolData(self,x.Symbol)
       for x in changes.RemovedSecurities:
            data = self.dataBySymbol.pop(x.Symbol,None)
            if data:
                data.dispose()
            
   def OnData(self, data):
        if self.IsWarmingUp:return
        
        # Loop through our equities
        for Symbol in self.dataBySymbol:
          symboldata = self.dataBySymbol[Symbol]
          if data.ContainsKey(Symbol) and data[Symbol] is not None and symboldata.IsReady() :
            # Set local variables
            close = data.Bars[Symbol].Close
            quantity = self.CalculateOrderQuantity(Symbol,1*0.5)
         
            # Define spread
            AskPrice = self.Securities[Symbol].AskPrice
            BidPrice = self.Securities[Symbol].BidPrice
            Spread = (AskPrice - BidPrice)
            
            
            # Create RSI variables 
            RSI1 = symboldata.rsi1.Current.Value
            RSI2 = symboldata.rsi2.Current.Value
        
            # Setup Open and Close Prices and Bars
            symboldata.window.Add(data[Symbol])
            if not (symboldata.window.IsReady and symboldata.daily.IsReady): return
            previous_bar_close = symboldata.window[0].Close
            previous_day_close = symboldata.daily[0].Close
            
            # Calculate change from close
            change_from_close = ((previous_bar_close / previous_day_close) -1) 
      
            #Obtain Low of Day and Update 
            bar = data[Symbol]
            
            if not bar.IsFillForward and symboldata.lowestPrice  < 0:
                symboldata.openPrice = bar.Open
                symboldata.lowestPrice = bar.Low
            
            if symboldata.lowestPrice < 0:
                return
            
            price = bar.Low
            if price < symboldata.lowestPrice:  # If we observe a new low
               symboldata.lowestPrice = price
         
            # IMPORTANT!!! Time variables to set open/close times and compare them to current time.
            
            # Convert times to variables (necessary for time comparison)
            currentTime = self.Time
            openTime = time(9,30)
            closeTime = time(12,0)
            
            # Convert string to format that can be compared (comparison does not work if you do not do this)
            # These comparisons are to test it works, before implementing them in the # Buy Conditions function below
            # It is OK to comment them out here, as they are just for testing. Real function below.
            #currentTime.strftime('%H%M') >= openTime.strftime('%H%M')
            #currentTime.strftime('%H%M') <= closeTime.strftime('%H%M')
            
            # Buy Conditions 
            if not self.Portfolio[Symbol].Invested and symboldata.ticket is None and (currentTime.strftime('%H%M') >= openTime.strftime('%H%M') and currentTime.strftime('%H%M') <= closeTime.strftime('%H%M')):
                self.current = self.Time
               
                # If buy conditions are satisfied then place MarketOrder
                if  ((RSI1 <=70 and RSI1 >=20) and (RSI2 >0 and RSI2 <=25) and (Spread >=0 and Spread <=0.01) and (change_from_close >= -0.11 and change_from_close <= -0.04) and (previous_bar_close <= symboldata.lowestPrice)): 
                    symboldata.ticket= self.MarketOrder(Symbol, quantity, tag ="Market Buy") and self.Debug(f"RSI1 Value: {RSI1}, RSI2 Value: {RSI2}, Change from close: {change_from_close}, Spread: {Spread}, Prev Bar Close: {previous_bar_close}, Time: {self.Time}") 
    
                    # Place Profit take and Stop Loss orders then reset to None
                    self.LimitOrder(Symbol, -quantity, close * 1.03, tag = "Profit Take")
                    self.StopMarketOrder(Symbol, -quantity, close * 0.99, tag = "Stopped Out")
                    symboldata.ticket = None
            
            # Close position if open for more than 15 minutes and set ticket to None
            if  (self.Time - self.current).seconds == 900:
                    self.Liquidate(Symbol)
                    symboldata.ticket = None
                    
   # Cancel remaining order if limit order or stop loss order is executed 
   def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        
        if order.Status == OrderStatus.Filled:
            if order.Type == OrderType.Limit or order.Type == OrderType.StopMarket:
                self.Transactions.CancelOpenOrders(order.Symbol)
                
        if order.Status == OrderStatus.Canceled:
            self.Log(str(orderEvent))
            
   #Reset Daily :Pointer
   def OnEndOfDay(self, symbol):
       for Symbol in self.dataBySymbol:
         self.dataBySymbol[Symbol].lowestPrice = -1