import numpy as np
from datetime import datetime
import pandas as pd

class AppendingInvestedToCoasre(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2018, 6, 20)   
        self.SetEndDate(2018, 9, 1)     
        self.coarse_amount = 5     
        self.SetCash(25000)           
        self.AddUniverse(self.CoarseSelectionFunction)         
        self.UniverseSettings.Resolution = Resolution.Daily    
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) 
        self.Data= {}
        self.stocks = []

    def CoarseSelectionFunction(self, coarse): 
        sortedByDollarVolume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)[:20]
        coarse_stocks = [x.Symbol for x in sortedByDollarVolume if x.Price > 15 and x.DollarVolume >= 900000 and x.HasFundamentalData == True][:self.coarse_amount]
        invested = [ x.Symbol for x in self.Portfolio.Values if x.Invested ]
        self.stocks =  np.unique(invested + coarse_stocks)
        
        self.Plot('Symbols', 'coarse_stocks', len(coarse_stocks))      
        self.Plot('Symbols', 'invested', len(invested))     
        self.Plot('Symbols', 'stocks', len(self.stocks))
        
        return self.stocks

        
    def OnSecuritiesChanged(self, changes):

        for security in changes.AddedSecurities:

            symbol = security.Symbol
            if symbol not in self.Data:
                self.Data[symbol] = SymbolData(self, symbol)

        for security in changes.RemovedSecurities:
            symbol = security.Symbol

            if symbol in self.Data:
                symbolData = self.Data.pop(symbol, None)
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.consolidator) 
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.consolidatorMinute)

    # =============== ON DATA ============================ #
    
    def OnData(self, data):
        selected = [] 
        symbols_in_onData = []
        
        for symbol in self.Data.keys():
            symbols_in_onData.append(symbol.Value)
            symbolData = self.Data[symbol]

            b1 = symbolData.Bars[0].Close > symbolData.Bars[1].Close
            b2 = symbolData.Bars[1].Close > symbolData.Bars[2].Close
            if b1 and b2:
                selected.append(symbol)
                
        for sec in self.Portfolio.Keys:
            if sec not in selected:
                self.SetHoldings(sec, 0)
                
        for sec in selected:       
            self.SetHoldings(sec, 0.95/len(selected))    
            
            

#============= SYMBOL DATA CLASS ========================== #        
class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        
        #self.emaOne = algorithm.EMA(symbol, 1, Resolution.Minute)
        self.ema3 = algorithm.EMA(symbol, 3, Resolution.Daily)
        self.ema4 = algorithm.EMA(symbol, 4, Resolution.Daily)
        self.ema50 = algorithm.EMA(symbol, 50, Resolution.Daily)
        self.ema200 = algorithm.EMA(symbol, 200, Resolution.Daily)
        self.ema200D = algorithm.EMA(symbol, 200, Resolution.Daily)
        self.ema200Window = RollingWindow[float](20)
        
        self.bb = algorithm.BB(symbol, 20, 2, MovingAverageType.Simple, Resolution.Daily)
        self.bbWindow = RollingWindow[float](5)
        self.bb17 = algorithm.BB(symbol, 20, 1.5, MovingAverageType.Simple, Resolution.Daily)
        self.bb17Window = RollingWindow[float](5)
        self.bb17LowerBandWindow = RollingWindow[float](5)
        
        self.momp = algorithm.MOMP(symbol, 30, Resolution.Daily)
        self.mompWindow = RollingWindow[float](200)
        
        self.rsi = algorithm.RSI(symbol, 14 , Resolution.Daily)
        self.rsiWindow = RollingWindow[float](20)
        
        self.macd = {}  #do I need?
        self.macd = algorithm.MACD(symbol, 12, 26, 9, MovingAverageType.Exponential, Resolution.Daily)
        self.macdWindow = RollingWindow[float](5)
        self.macdHistogramWindow = RollingWindow[float](5)
        
        #====== V-MACD =============
        self.vwap12 = algorithm.VWAP(symbol, 12, Resolution.Daily)   #12 period VWAP
        self.vwap26 = algorithm.VWAP(symbol, 26, Resolution.Daily)   #26 perios VWAP
        self.vmacd = IndicatorExtensions.Minus(self.vwap12, self.vwap26) #vwap26 - vwap12
        self.vmacdSignal = IndicatorExtensions.EMA(self.vmacd, 9)   
        self.vmacdHistogram = IndicatorExtensions.Minus(self.vmacd, self.vmacdSignal) #swap
        self.vmacdHistogramWindow = RollingWindow[float](5)
        
        self.close_window = RollingWindow[float](22)
        
        self.Bars = RollingWindow[IBaseDataBar](22) # Rolling window for data bars
        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
    
        self.BarsWeek = RollingWindow[IBaseDataBar](22) # Rolling window for data bars
        self.consolidatorWeek = TradeBarConsolidator(timedelta(days=1))
        self.consolidatorWeek.DataConsolidated += self.OnDataConsolidatedWeek
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidatorWeek)
        
        self.BarsMinute = RollingWindow[IBaseDataBar](22) # Rolling window for data bars
        self.consolidatorMinute = TradeBarConsolidator(timedelta(days=2))
        self.consolidatorMinute.DataConsolidated += self.OnDataConsolidatedMinute
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidatorMinute)
        
        #=====History Calls========
    
        history = algorithm.History(self.symbol, 25, Resolution.Daily).loc[self.symbol]
        for time, row in history.iterrows():
            tradebar = TradeBar(time, self.symbol, row.open, row.high, row.low, row.close, row.volume)
            self.consolidatorMinute.Update(tradebar)
    
        history = algorithm.History(self.symbol, 200, Resolution.Daily).loc[self.symbol]
        for time, row in history.iterrows():
            tradebar = TradeBar(time, self.symbol, row.open, row.high, row.low, row.close, row.volume)
            self.consolidator.Update(tradebar)
            self.consolidatorWeek.Update(tradebar)
            self.vwap12.Update(tradebar)
            self.vwap26.Update(tradebar)
        
    # Warm up indicators
        history = algorithm.History([symbol], 201, Resolution.Daily)
        for time, row in history.loc[symbol].iterrows():
            #self.emaOne.Update(time, row["close"])
            self.ema200.Update(time, row["close"])
            self.momp.Update(time, row["close"])
            self.ema3.Update(time, row["close"])
            self.ema4.Update(time, row["close"])
            
            self.bb.Update(time, row["close"])
            self.bb17.Update(time, row["close"])

            
            self.rsi.Update(time, row["close"])
            self.macd.Update(time, row["close"])
            
            self.close_window.Add(row["close"]) #Per anwser it needs the close only
            
            #Warm the Rolling Windows
            if self.ema200.IsReady:
                self.ema200Window.Add(self.ema200.Current.Value)
            
            if self.rsi.IsReady:
                self.rsiWindow.Add(self.rsi.Current.Value)
                
            if self.bb.IsReady:
                self.bbWindow.Add(self.bb.Current.Value)
            if self.bb17.IsReady:
                self.bb17Window.Add(self.bb17.Current.Value)
                self.bb17LowerBandWindow.Add(self.bb17.LowerBand.Current.Value)
                
            if self.macd.IsReady:
                self.macdWindow.Add(self.macd.Current.Value)
                self.macdHistogramWindow.Add(self.macd.Histogram.Current.Value)
                self.vmacdHistogramWindow.Add(self.vmacdHistogram.Current.Value)
                
            if self.momp.IsReady:
                self.mompWindow.Add(self.momp.Current.Value)
                
        history = algorithm.History([symbol], 200, Resolution.Daily)
        for time, row in history.loc[symbol].iterrows():
            self.ema200D.Update(time, row["close"])
            
    # Consolidators            
    def OnDataConsolidated(self, sender, bar):
        self.Bars.Add(bar)
        
    def OnDataConsolidatedMinute(self, sender, bar):
        self.BarsMinute.Add(bar)
        
    def OnDataConsolidatedWeek(self, sender, bar):
        self.BarsWeek.Add(bar)
        
    @property
    def IsReady(self):
        return self.Bars.IsReady and self.BarsWeek.IsReady and self.rsi.IsReady and self.ema200.IsReady and self.macd.IsReady #and self.macdWindow.IsReady