import numpy as np
import pandas as pd


class TrendfollowingEffectsInStocks(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2014, 1, 1)
        self.SetEndDate(2019,1,1)
        self.SetCash(1000000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        self.symbols = []
        #Store the maximum prices
        self.history = pd.DataFrame()
        self.symbolData = {}
        self.month = 0

    def CoarseSelectionFunction(self, coarse):
        # Avoid missing data as it happens on Oct 18th 2018
        if len(coarse) == 0:
            self.Log(f'{self.Time}: No data for coarse!')
            return self.symbols
        # update universe once every month
        if self.Time.month == self.month:
            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[:30]]
        self.month = self.Time.month
        return self.symbols
        
        
    def OnData(self, data):
        # If no holdings, scan the universe to see if any stock triggers buying signal
        self.Log(f'{self.Time} OnData begins')
        
        if not self.Portfolio.Invested:
            self.Log(f'{self.Time} not invested')
            long_list = []
            for symbol in self.symbols:
                if data.ContainsKey(symbol) and data[symbol] is not None and symbol in self.symbolData:
                    self.symbolData[symbol].Prices.Add(data[symbol].Close)
                    if data[symbol].Close > self.symbolData[symbol].High:
                        self.symbolData[symbol].High = data[symbol].Close
                        long_list.append(symbol)
                        
            
            for symbol in long_list:
                self.SetHoldings(symbol, 1/len(long_list))
                #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
                self.StopMarketOrder(symbol, -self.Portfolio[symbol].Quantity , self.symbolData[symbol].Prices[1] - 4.5*self.symbolData[symbol].ATR.Current.Value)
            

        # If there are holdings in portfolio, then stop adding new stocks into long_list. We rebalance the portfolio, and watch all the holdings exit   
        else:
            #Take care of the existing holdings and place StopMarketOrder to try to exit
            for symbol in self.symbols:
                if symbol in self.Portfolio.Keys:
                    self.Log(f'{self.Time} {str(symbol)} in the portfolio')
                    holding = self.Portfolio[symbol]
                    if holding.Invested and data.ContainsKey(symbol) and data[symbol] is not None and symbol in self.symbolData:
                        # If price goes up then also adjust the stopPrice up, else keep the stopPrice constant, this is so called trailing stop loss
                        # Here, to avoid look-ahead bias, we can only know yesterday's price when placing stop loss orders
                        if self.symbolData[symbol].Prices[1] > self.symbolData[symbol].Prices[2]:
                            self.StopMarketOrder(symbol, -holding.Quantity , self.symbolData[symbol].Prices[1] - 4.5*self.symbolData[symbol].ATR.Current.Value)
                        else:
                            self.StopMarketOrder(symbol, -holding.Quantity , self.symbolData[symbol].Prices[2] - 4.5*self.symbolData[symbol].ATR.Current.Value)
            
            # Rebalance the holdings to equally weighted
            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)

            self.symbolData[symbol].ATR.Reset()
            #Set the high
            self.symbolData[symbol].High = high[ticker]
            
            for value in close[ticker]:
                self.symbolData[symbol].Prices.Add(value)
            #Set the ATR indicator
            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.Symbol = symbol
        self.ATR = algorithm.ATR(symbol, 10)
        self.High = 0
        #Track the recent 3-days' prices so that we can build trailing stop loss signal
        self.Prices = RollingWindow[float](3)