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, 3, 1)  # Set Start Date\
        self.SetEndDate(2021, 5, 1)
       
        # Holds all of our data keyed by each symbol
        self.Data = {}
        # Contains all of our equity symbols
        self.UniverseSettings.Resolution = Resolution.Minute

        self.UniverseSettings.SetDataNormalizationMode = DataNormalizationMode.Raw
        
        self.AddUniverseSelection(ScheduledUniverseSelectionModel(self.DateRules.Every(DayOfWeek.Monday),
        self.TimeRules.At(0,0), self.CoarseSelectionFilter(self)))
        
        self.__numberOfSymbols = 10

        #self.Schedule.On(self.DateRules.On(2021, 3, 1), self.TimeRules.At(0, 0), self.CoarseSelectionFilter(self))

        self.AddUniverse(self.CoarseSelectionFilter)
        
        self.symbol = Symbol.Create("TSLA", SecurityType.Equity, Market.USA)

        self.UniverseSettings.Leverage = 2
        
    
    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):
        
        # Ver estos cambios
        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.09)
           
            
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