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

class AppendingInvestedToCoasre(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2018, 6, 20)   # Start date of backtest
        self.SetEndDate(2018, 9, 1)     # End date of backtest
        self.coarse_amount = 5     
        self.SetCash(25000)            # Amount of cash in account for backtest
        self.AddUniverse(self.CoarseSelectionFunction)          # Adding CoarseSelectionFunction as Universe
        self.UniverseSettings.Resolution = Resolution.Daily    # Sets Universe resolution to minute
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) 
        self.Data= {}

    def CoarseSelectionFunction(self, coarse): 
        sortedByDollarVolume = sorted(coarse, key=lambda c: c.DollarVolume, reverse=True)
        stocks = [x.Symbol for x in sortedByDollarVolume if x.Price > 15 and x.DollarVolume >= 900000 and x.HasFundamentalData == True] #HasFundamentalData effectivley removes indexs/ETFs
        invested = [ x.Symbol for x in self.Portfolio.Values if x.Invested ]
        invested_tickers = [ x.Symbol.Value for x in self.Portfolio.Values if x.Invested ]
        
        #Debug Code
        #self.Debug(f"(Coarse) Invested in {len(invested)} stocks: {invested_tickers}")

        stocks =  invested + stocks
        
        try:
            if invested[0] in stocks[:self.coarse_amount]:
                index_of_invested = stocks.index(invested[0])
                #self.Debug(f"The index of invested[0] in stocks is {index_of_invested} and is {stocks[index_of_invested]}")
        except:
            pass
        try:
            if invested[1] in stocks[:self.coarse_amount]:
                index_of_invested = stocks.index(invested[1])
                #self.Debug(f"The index of invested[1] in stocks is {index_of_invested} and is {stocks[index_of_invested]}")
        except:
            pass

        return_amount = self.coarse_amount + len(invested)
        #self.Debug(f"We are returning {(return_amount)} symbols from the Coarse function")
        self.Debug(f"Coarse - returns: {[stock.Value for stock in stocks[:return_amount]]}")
        return stocks[:return_amount]

        
    def OnSecuritiesChanged(self, changes):
        i=0
        syms=[]
        for security in changes.AddedSecurities:
            i+=1
            syms.append(security.Symbol.Value)
            symbol = security.Symbol
            if symbol not in self.Data:
                self.Data[symbol] = SymbolData(self, symbol)
        #self.Debug(f"in Added Securities we looped through {i} added securities: {syms}")
        self.Debug(f"OnSecuritiesChanged - Added securities: {[security.Symbol.Value for security in changes.AddedSecurities]}")
        
        i=0
        syms=[]
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            i+=1
            syms.append(security.Symbol.Value)
            if symbol in self.Data:
                symbolData = self.Data.pop(symbol, None)
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.consolidator) # Remove subcription for symbols removed from universe
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.consolidatorMinute)
        #self.Debug(f"in Removed Securities we looped through {i} removed securities: {syms}") 
        self.Debug(f"OnSecuritiesChanged - Removed securities: {[security.Symbol.Value for security in changes.RemovedSecurities]}")

    # =============== ON DATA ============================ #
    def OnData(self, data):
        #self.Debug(f"===On Data===")
        self.Debug(f"OnData - symbols passed: {[d.Value for d in data.Bars.keys()]}")

        invested = [ x.Symbol.Value for x in self.Portfolio.Values if x.Invested ]
        #self.Debug(f"\n(onData) Invested in {len(invested)} stocks")
        #if len(invested) > 0:
            #for i in range(0,len(invested)):
                #self.Debug(f"{i+1}/{len(invested)}: {invested[i]}")
        
        symbols_in_onData = []
        for symbol in self.Data.keys():
            symbols_in_onData.append(symbol.Value)
            #if symbol in invested:
                #self.Debug(f"{symbol.value} was in invested")

            symbolData = self.Data[symbol]
            #if not symbolData.IsReady: 
                #self.Debug(f"{symbol.Value} was not ready")
            
            #Check for sell conditions
            try:
                if symbol in invested:
                    s1 = symbolData.Bars[0].Close < symbolData.Bars[1].Close
                    if s1:
                        self.Liquidate(symbol)
                        self.Debug(f"----- OnData - {symbol.Value} was liquidated -----")
                else:
                    pass
            except:
                self.Debug(f"!!!!! OnData - ERROR in selling code: with {symbol.Value} !!!!!")
                
                
            #Check for Buy conditions
            try:
                b1 = symbolData.Bars[0].Close > symbolData.Bars[1].Close
                b2 = symbolData.Bars[1].Close > symbolData.Bars[2].Close
                if b1 and b2:
                    self.SetHoldings(symbol, 0.5)
                    self.Debug(f"+++++ OnData - Conditions to buy met for {symbol.Value} +++++")
                else:
                    pass
            except:
                self.Debug(f"!!!!! OnData - ERROR in buy code: with {symbol.Value} !!!!!")
        
        #len_onData = len(symbols_in_onData)        
        #self.Debug(f" {len_onData} symbols were looped through onData slice {symbols_in_onData}")
        #i=0
        #for symbol in invested:
            #if symbol in symbols_in_onData:
                #i+=1
                #self.Debug(f" {symbol} found in onData symbols")
        #self.Debug(f"{i}/{len(invested)} invested symbols found in the {len_onData} symbols)")
            
  
  

  
  
  
    #======= ORDER EVENTS===================
    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)
                self.Debug(f"----- OnOrderEvent - order cancelled: {oder.Symbol.Value} -----")
                
        if order.Status == OrderStatus.Canceled:
            pass
            #self.Log(str(orderEvent))    
            

#============= 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