from System import *
from QuantConnect import *
from QuantConnect.Data.Consolidators import *
from QuantConnect.Data.Market import *
from QuantConnect.Orders import OrderStatus
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Indicators import *
import numpy as np
from datetime import timedelta, datetime

class MultipleSymbolConsolidationAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        
        self.SetStartDate(2021, 1, 26)  # Set Start Date\
        self.SetEndDate(2021, 4, 26)
        # Holds all of our data keyed by each symbol
        self.Data = {}
         
        # Contains all of our equity symbols

        self.AddUniverse(self.CoarseSelectionFilter)
        self.UniverseSettings.Resolution = Resolution.Minute   
       
        self.UniverseSettings.SetDataNormalizationMode = DataNormalizationMode.Raw
        self.UniverseSettings.Leverage = 1 
        
        self.__numberOfSymbols = 10
        
    
    def CoarseSelectionFilter(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda c: c.Volume, reverse=True)
        filteredByPrice = [c.Symbol for c in sortedByDollarVolume if c.Price > 10 ]
        return filteredByPrice[:self.__numberOfSymbols]
            
            
    def OnSecuritiesChanged(self, changes):
        
        # This is the period of bars we'll be creating
        BarPeriod = TimeSpan.FromMinutes(10)
        
        # This is the period of our sma indicators
        SimpleMovingAveragePeriod = 20
        
        # This is the number of consolidated bars we'll hold in symbol data for reference
        RollingWindowSize = 10   
        
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.Data[symbol] = SymbolData(self, symbol, BarPeriod, SimpleMovingAveragePeriod, RollingWindowSize)
                
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.Data:
                self.SubscriptionManager.RemoveConsolidator(symbol, self.Data[symbol].consolidator)
                symbolData = self.Data.pop(symbol, None)
                if security.Invested:
                    self.Liquidate(symbol, "Universe Removed Security")

    def OnDataConsolidated(self, sender, bar):
        if bar.Symbol in self.Data:
            self.Data[bar.Symbol].SMA.Update(bar.EndTime, bar.Close)
            if self.Data[bar.Symbol].SMA.IsReady:
                sma = self.Data[bar.Symbol].SMA.Current.Value
                self.Data[bar.Symbol].HistoricalSMA.Add(sma)
            self.Data[bar.Symbol].Bars.Add(bar)

    # OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
    # Argument "data": Slice object, dictionary object with your stock data 
    def OnData(self,data):
        # loop through each symbol in our structure
        for symbol, symbolData in self.Data.items():
            # this check proves that this symbol was JUST updated prior to this OnData function being called
            if symbolData.IsReady() and symbolData.WasJustUpdated(self.Time):
                
                if not self.Portfolio[symbol].Invested:
                    self.SetHoldings(symbol, 0.1)
                
                '''------------------------- Trading logic  ----------------------------------- '''
                a = symbolData.Bars[0].Close
                b = symbolData.HistoricalSMA[1]     # second most recent item
                
                if a >= b:
                    self.MarketOrder(symbol, 1) 
                    continue
            
                if a <= b:  
                    self.MarketOrder(symbol, -1) 
                    continue
                
class SymbolData(object):
    
    def __init__(self, algo, symbol, barPeriod, smaPeriod, windowSize):
        self.Symbol = symbol
        
        # The period used when population the Bars rolling window
        self.BarPeriod = barPeriod
        self.Bars = RollingWindow[IBaseDataBar](windowSize)
        
        # The simple moving average indicator for our symbol
        self.SMA = SimpleMovingAverage(algo.CreateIndicatorName(symbol, "SMA" + str(smaPeriod), Resolution.Minute), smaPeriod)
        
        # define a consolidator to consolidate data for this symbol on the requested period
        self.consolidator = TradeBarConsolidator(barPeriod) 
        # write up our consolidator to update the indicator
        self.consolidator.DataConsolidated += algo.OnDataConsolidated
        # we need to add this consolidator so it gets auto updates
        algo.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
        sma_lookback = 10
        self.HistoricalSMA = RollingWindow[float](sma_lookback)
        
        
    # Returns true if all the data in this instance is ready (indicators, rolling windows, ect...)
    def IsReady(self):
        return self.Bars.IsReady and self.SMA.IsReady and self.HistoricalSMA.IsReady

    # Returns true if the most recent trade bar time matches the current time minus the bar's period, this
    # indicates that update was just called on this instance
    def WasJustUpdated(self, current):
        return self.Bars.Count > 0 and self.Bars[0].Time == current - self.BarPeriod