import pandas as pd
import numpy as np
from datetime import datetime, timedelta
from QuantConnect.Data.UniverseSelection import *
from scipy.stats import linregress

class Momentum(QCAlgorithm):

    def __init__(self):
    # set the flag for rebalance
        self.reb = 1
    # Number of stocks to pass CoarseSelection process
        self.num_coarse = 500
    # Number of stocks to long
        self.num_fine = 10
        self.symbols = None
        
    def Initialize(self):
        self.SetStartDate(2018, 8, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # set the flag for rebalance
        self.reb = 1
        self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
        # for market state filter
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.Schedule.On(self.DateRules.MonthStart(self.spy), 
        self.TimeRules.AfterMarketOpen(self.spy,5), Action(self.rebalance))
        
    def CoarseSelectionFunction(self, coarse):
    # if the rebalance flag is not 1, return null list to save time.
        if self.reb != 1:
            return self.long 
        
        # make universe selection once a month
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)  
        filtered = [x.Symbol for x in sortedByDollarVolume if x.HasFundamentalData]
        
        # filtered down to the 500 most liquid stocks
        return filtered[:self.num_coarse]
        
    def FineSelectionFunction(self, fine):
    # return null list if it's not time to rebalance
        if self.reb != 1:
            return self.long 
            
        # drop counter (will update back to 1 after rebalancement has occurred)
        self.reb = 0
        
        # create dictionaries to store the indicator values
        
        stock_filter1 = {}
        
        # sort the fine list by their momentum 
        for security in fine: 
            
            hist = self.History(security.Symbol, timedelta(days=365), Resolution.Daily)
            
            if "close" in hist.columns and len(hist) > 239: 
                hist["log"] = np.log(hist["close"])
                x1 = np.arange(hist["close"].count())
                slope, _, rvalue, _, _ = linregress(x1, hist["log"])
                coeff = slope*252*(rvalue**2)
                        
                # we now have a dictionary storing the values 
                stock_filter1[security.Symbol] = coeff
                
        # we only want the highest values for the coeff
        self.sorted1 = sorted(stock_filter1.items(), key=lambda d:d[1],reverse=True)
        sorted1_symbol = [x[0] for x in self.sorted1]
        
        
        # long the top 10
        self.long = sorted1_symbol[:self.num_fine]
        return self.long

    def OnData(self, data):
        pass
    
    def rebalance(self):
        # at the start of rebalancement, if the stock is no longer in the long list, liquidate
        long_list = self.long
        for i in self.Portfolio.Values:
            if (i.Invested):
                self.Liquidate(i.Symbol)
            
    # Assign each stock equally. Alternatively you can design your own portfolio construction method
        for i in self.long:
            self.SetHoldings(i, 1/self.num_fine)
        
        self.reb = 1