# https://quantpedia.com/Screener/Details/14
from QuantConnect.Indicators import *
class MomentumEffectAlgorithm(QCAlgorithm):
# Trading time
# The other side selling: not great when the market is always increasing
# 1) change the weight of selling 2) correlated with the market
# Setting up the threshold
# rsi index？or second derivative 
# short using bad stocks 
# restructure code (check)
# threshold: could be solved by sorting stocks?
# change stocks

    def Initialize(self):
        #self.AddEquity(self.temp_ticker, Resolution.Daily)
        
        self.SetStartDate(2015, 1, 1)  # Set Start Date
        self.SetEndDate(2021, 1, 1)    # Set Start Date       
        self.SetCash(100000)           # Set Strategy Cash
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.hparams = {'lookback': 1 * 7}
        self.num_coarse = 100
        self.num_fine = 50
        
        # Momentum-related variables
        self.moment = {}
        self.last_price = {}
        self.lookback = 1 * 7     # Momentum indicator lookback period
        self.last_trade_mom = 0
        self.volume = {}
        self.seen_day = {}
        
        # Long-related variables
        self.long_enable = True
        self.curr_long_selected = []
        self.selected_stocks_long = []
        self.num_long = 10
        
        # Short-related variables
        self.short_enable = False
        self.curr_short_selected = []
        self.selected_stocks_short = []
        self.num_short = 10
        
    def CoarseSelectionFunction(self, coarse):
        '''Drop securities which have no fundamental data or have too low prices.
        Select those with highest by dollar volume'''
        if self.Time.weekday() != 3:
            return Universe.Unchanged
        
        self.rebalance = True
        #self.month = self.Time.month
        self.day = self.Time.day
        
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 10],
            key=lambda x: x.DollarVolume, reverse=True)

        return [x.Symbol for x in selected[:self.num_coarse]]

    def FineSelectionFunction(self, fine):
        '''Select security with highest market cap for long'''

        fine = [f for f in fine if f.ValuationRatios.PERatio > 0
                               and f.EarningReports.BasicEPS.TwelveMonths > 0
                               and f.EarningReports.BasicAverageShares.ThreeMonths > 0]

        selected_long = sorted(fine,
            key=lambda f: f.MarketCap,
            reverse=True)
        selected_short = sorted(fine,
            key=lambda f: f.MarketCap,
            reverse=False)
        self.selected_stocks_long = [x.Symbol for x in selected_long[:self.num_fine]]
        self.selected_stocks_short = [x.Symbol for x in selected_short[:self.num_fine]]
        return self.selected_stocks_long + self.selected_stocks_short
    
    def OnSecuritiesChanged(self, changes):
        # Clean up data for removed securities and Liquidate
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if self.moment.pop(symbol, None) is not None:
                self.Liquidate(symbol, 'Removed from universe')

        for security in changes.AddedSecurities:
            if security.Symbol not in self.moment:
                self.prepare_indicator(security.Symbol)
                    
    def prepare_indicator(self, ticker):
        history = self.History([ticker], self.hparams['lookback'] + 1, Resolution.Daily)
        for time, row in history.loc[ticker].iterrows():
            if ticker not in self.last_price:
                self.last_price[ticker] = row["close"]
                self.volume[ticker] = row["volume"]
                self.seen_day[ticker] = 1
            elif ticker not in self.moment:
                self.moment[ticker] = (row["close"] - self.last_price[ticker]) / self.last_price[ticker]
                self.volume[ticker] = self.volume[ticker] + row["volume"]
                self.seen_day[ticker] += 1
            else:
                mu = row["volume"] * self.seen_day[ticker] / self.volume[ticker] 
                percentage_return = (row["close"] - self.last_price[ticker]) / self.last_price[ticker]
                self.moment[ticker] =  (self.moment[ticker] + mu * percentage_return) / (1 + mu)
                self.volume[ticker] = self.volume[ticker] + row["volume"]
                self.seen_day[ticker] += 1
                self.Debug("Mu: ")
                self.Debug(mu)
            if ticker in self.last_price and ticker in self.moment:
                self.Debug('return: ')
                self.Debug(row["close"] - self.last_price[ticker])
                self.Debug('momentum: ')
                #self.Debug(self.moment[self.temp_ticker])
            self.last_price[ticker] = row["close"]
    
    def update_indicator(self, ticker):
        #guranteed to have moment, volume, seen_day exist for 'ticker'
        if ticker not in self.moment:
            self.prepare_indicator(ticker)
        if self.Securities[ticker].Close == 0:
            return
        percentage_return = (self.Securities[ticker].Close  - self.last_price[ticker]) / self.last_price[ticker]
        mu = self.Securities[ticker].Volume * self.seen_day[ticker] / self.volume[ticker]
        self.seen_day[ticker] += 1
        self.volume[ticker] += self.Securities[ticker].Volume
        self.moment[ticker] = (self.moment[ticker] + mu * percentage_return) / (1 + mu)
        self.last_price[ticker] = self.Securities[ticker].Close

    def OnData(self, data):
        self.Debug(str(self.Time.month))
        #self.Debug(self.moment[self.temp_ticker])
        
        #if self.Time.weekday() != 2 and self.Time.weekday() != 4:
        #    return
        
        for symbol, _ in self.moment.items():
            self.update_indicator(symbol)
        
        if self.long_enable:
            mom_long = [(k, v) for k, v in self.moment.items() if k in self.selected_stocks_long]
            sorted_mom_long = sorted([k for k,v in mom_long],
                key=lambda x: self.moment[x], reverse=True)
            selected_long = sorted_mom_long[:self.num_long]

            # Liquidate securities that are not in the list
            for symbol in self.curr_long_selected:
                if symbol not in selected_long:
                    self.Liquidate(symbol, 'Not selected')
            
            # Buy selected securities
            for symbol in selected_long:
                self.SetHoldings(symbol, 0.8 / self.num_long)
        if self.short_enable:
            mom_short = [(k, v) for k, v in self.moment.items() if k in self.selected_stocks_short]
            sorted_mom_short = sorted([k for k,v in mom_short],
                key=lambda x: self.moment[x], reverse=False)
            selected_short = sorted_mom_short[:self.num_short]

            # Liquidate securities that are not in the list
            for symbol in self.curr_short_selected:
                if symbol not in selected_short:
                    self.Liquidate(symbol, 'Not selected')
            
            # Buy selected securities
            for symbol in selected_short:
                self.SetHoldings(symbol, -0.8 / self.num_short)