import numpy as np
import pandas as pd


class TrendfollowingEffectsInStocks(QCAlgorithm):
    def Initialize(self):
        
        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(2018,1,1)
        self.SetCash(1000000)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        
        self.month = 0
        self.count = 30
        self.symbols = []
        self.long_list = []
        # Store the data and indicators for each symbol
        self.symbolData = {}

    def CoarseSelectionFunction(self, coarse):
        
        # Update universe once every month
        if self.Time.month == self.month:
            return self.symbols
        
        self.month = self.Time.month
        
        if len(coarse) == 0:
            self.Log(f'{self.Time}: No data for coarse!')
            return self.symbols
        
        selected = [x for x in coarse if x.HasFundamentalData and x.Price > 3]
        filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True)
        
        self.symbols = [x.Symbol for x in filtered[:self.count]]
        return self.symbols
        
        
    def OnData(self, data):

        for symbol in self.symbols:
           if data.ContainsKey(symbol) and data[symbol] is not None:
                close = data[symbol].Close
                symbolData = self.symbolData[symbol]
                symbolData.Prices.Add(close)
                
                if close > symbolData.High:
                    symbolData.High = close
                    self.long_list.append(symbol)
            
            
        for symbol in self.symbols:
            if self.Portfolio.ContainsKey(symbol):
                if self.Portfolio[symbol].Invested and symbol in self.long_list:
                    self.long_list.remove(symbol)
                    
        if len(self.long_list) == 0:
            return
        else:
            for symbol in self.long_list:
                self.SetHoldings(symbol, 1/len(self.long_list))
                    
       
        
        # If there are holdings in portfolio, then stop adding new stocks into long_list.
        # We rebalance the portfolio, and watch all the holdings exit   
        # Take care of the existing holdings and place StopMarketOrder to try to exit
        for symbol in self.symbols:
            if self.Portfolio.ContainsKey(symbol):
                if self.Portfolio[symbol].Invested:
                    self.Log(f'{self.Time}: portfolio has {self.Portfolio[symbol].Quantity} of {str(symbol)}')

                    symbolData = self.symbolData[symbol]
                    if symbolData.StopLoss is None:
                        # To avoid look-ahead bias, we can't use today's close price to place a stop market order because
                        # when we place the order, we wouldn't know that day's close price. So use yesterday's price
                        symbolData.StopPrice = symbolData.Prices[1] - 3.5*symbolData.ATR.Current.Value
                        symbolData.StopLoss = self.StopMarketOrder(symbol, -self.Portfolio[symbol].Quantity, symbolData.StopPrice)
                        self.Log(f'{self.Time}: placed {-self.Portfolio[symbol].Quantity} market stop orders at price {symbolData.StopPrice}')
                    else:
                        self.Transactions.CancelOpenOrders(symbol)
                        
                        # Only update submitted orders
                        if symbolData.StopLoss.Status != OrderStatus.Submitted:
                            continue
                        
                        updateOrderFields = UpdateOrderFields()
                        
                        # If price goes up (price at t-1 > price at t-2) then also adjust the stopPrice up, else keep the stopPrice constant, this is so-called trailing stop loss
                        # Here, we only know yesterday's price when placing stop loss orders. To avoid look-ahead bias, we don't use today's close
                        updateOrderFields.StopPrice = max(symbolData.Prices[1]- symbolData.ATR.Current.Value, symbolData.Prices[2]- symbolData.ATR.Current.Value)
                        self.Log(f"{symbolData.StopLoss} : {updateOrderFields} : {symbolData.Prices}")
                        symbolData.StopLoss.Update(updateOrderFields)


        # Rebalance the holdings to equally weighted
        if sum(x.Invested for x in self.Portfolio.Values) == 0:
            return
        else:
            target = 1 / sum(x.Invested for x in self.Portfolio.Values)
            for holding in self.Portfolio.Values:
                if holding.Invested:
                    self.SetHoldings(holding.Symbol, target)
    
        
    def OnSecuritiesChanged(self, changes):
        # Liquidate positions of removed securities
        # for security in changes.RemovedSecurities:
        #     if security.Invested:
        #         self.Liquidate(security.Symbol)
        
        #Call history request to initialize ATR indicator and prices rolling window
        history = self.History(self.symbols, 252*5, Resolution.Daily)
        high = history.high.unstack(level=0).max()
        close = history.close.unstack(level=0).tail(3).fillna(0)
    
        for symbol in self.symbols:
            ticker = str(symbol)
            if ticker not in close.columns or ticker not in high.index:
                self.symbols.remove(symbol)
                continue
            if symbol not in self.symbolData:
                self.symbolData[symbol] = SymbolData(self, symbol)

            #Set the high
            self.symbolData[symbol].High = high[ticker]
            
            for value in close[ticker]:
                self.symbolData[symbol].Prices.Add(value)
            
            # Warm up the ATR indicator
            self.symbolData[symbol].ATR.Reset()
            for index, row in history.loc[ticker].tail(10).iterrows():
                bar = TradeBar(index, symbol, row.open, row.high, row.low, row.close, row.volume)
                self.symbolData[symbol].ATR.Update(bar)
        

class SymbolData:
    def __init__(self, algorithm, symbol):
        self.ATR = algorithm.ATR(symbol, 10)
        
        # Track the security high
        self.High = 0
        
        # Track the recent 3-days' prices so that we can build trailing stop loss signal
        self.Prices = RollingWindow[float](3)
        
        # Save information for the Stop Loss order
        self.StopPrice = 0
        self.StopLoss = None