Overall Statistics
from QuantConnect.Data.UniverseSelection import *
import numpy as np
import pandas as pd
from math import ceil

### This is a basic monthly momentum rotation strategy

class MonthlyMomentumRotationAlgorithm(QCAlgorithm):
    '''Uses standard percent change over 200 days to determine momentum calculation'''

    def Initialize(self):
        '''While Quantopian sets cash and starting and ending dates in GUI, in QuantConnect we set them here.'''

        # Set dates and cash
        self.SetStartDate(2010, 1, 1)  #Set Start Date
        self.SetEndDate(2012, 7, 3)    #Set End Date
        self.SetCash(10000000)
        
        # Set settings to determine QC500 generator
        self.num_coarse = 1000
        self.num_fine = 500
        self.dollar_volume = {} # empty dictionary

        # Add the universe here, is similar to the Pipeline in Quantopian
        self.UniverseSettings.Resolution = Resolution.Daily # daily resolution is sufficient for what we need
        self.AddUniverse( self.CoarseSelectionFunction, self.FineSelectionFunction)

        # Add SPY data using AddEquity() method, will need it as a market gate
        # AddEquity() returns QuantConnect Equity object, Symbol property returns QuantConnect symbol object
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol # self.spy holds a symbol object

        # Scheduling functions
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(0, 0), Action(self.monthly_rebalance))
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY",1), Action(self.rebalance)) # carry out the rebalance 1 minute after market open
        
        # Set flags and initialize list
        self.rebalance_flag = 0
        self.first_month_trade_flag = 1
        self.symbols = None
        
        # Set settings for the algorithm
        self.num_stocks = 30 # hold a portfolio of 30 stocks
    
    # coarse selection and fine selection function are modelled on ConsituentsQC500GeneratorAlgorithm  https://github.com/QuantConnect/Lean/pull/1663
    # this can act as a proxy to S&P 500, as QuantConnect doesn't have an S&P 500.
    # coarse is a list of coarsefundamental type objects which have available properties:
    # https://www.quantconnect.com/docs/algorithm-reference/universes#Universes-Coarse-Universe-Selection

    def CoarseSelectionFunction(self, coarse):
        if self.rebalance_flag or self.first_month_trade_flag: # only want to run this if we have to rebalance, otherwise slows algorithm
            # x is a coarsefundamentaltype object
            filtered = [x for x in coarse if x.HasFundamentalData # boolean
                                          and x.Volume > 0 # this is a long
                                          and x.Price > 0] # this is a decimal
            # sort by output of lambda function (so DollarVolume property), reverse order, take top 1,000
            sort_filtered = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)[:self.num_coarse]
            for i in sort_filtered: # for each coarsefundamentaltype object
                self.dollar_volume[i.Symbol.Value] = i.DollarVolume # create key/value pair with key = string of symbol, and value the dollar volume, can use this dict later
            self.symbols = [x.Symbol for x in sort_filtered]  # list of symbol objects (not coarsefundamentaltype objects)
            return self.symbols
            # that is all we are going to do here, we will sort by top market cap (to approximate S&P 500 in fine selection)
        else:
            return self.symbols # previous symbols list
    
    # fine selection takes this further to a proxy for S&P 500, by using market cap and sector selection
    # more info on all the fine fundamental data here: https://www.quantconnect.com/docs/data-library/fundamentals#Fundamentals-Morningstar-US-Equity-Data
    def FineSelectionFunction(self, fine): # fine is a list of finefundamental type objects 
        if self.rebalance_flag or self.first_month_trade_flag:
            filtered_fine = [x for x in fine if (x.CompanyReference.CountryId == "USA") # 3 letter string
                                             and (x.CompanyReference.PrimaryExchangeID == "NYS" or x.CompanyReference.PrimaryExchangeID == "NAS") # string
                                             and ((self.Time - x.SecurityReference.IPODate).days > 180) # this is a timedate object
                                             and x.EarningReports.BasicAverageShares.ThreeMonths * (x.EarningReports.BasicEPS.TwelveMonths*x.ValuationRatios.PERatio) > 5e8]
                                             # calculate market cap using number of shares x EPS x PERatio (last 2 get the price per share, which is not available in fine selection)
            for i in filtered_fine: # for each element in filtered_fine list (which is a finefundamental type object)
                i.DollarVolume = self.dollar_volume[i.Symbol.Value] # using dictionary created in coarse selection, create a property of each finefundamental type object
                # which is the dollar volume. It appears we can create our own properties to finefundamental type objects
            percent = float(self.num_fine/len(filtered_fine)) # calculate how much we need to trim the list by (used in sector selection later)
            group_by_code = {}
            top_list = []
            for code in ["N", "M", "U", "T", "B", "I"]:
                # populate dictionary with code as key, CompanyReference.IndustryTemplateCode returns a string with any of these 6 options:
                # N=Normal (Manufacturing), M=Mining, U=Utility, T=Transportation, B=Bank, I=Insurance
                group_by_code[code] = list(filter(lambda x: x.CompanyReference.IndustryTemplateCode == code, filtered_fine))
                # filter (func in python) the list filtered_fine by industry template code, creates a list of elements for which lambda func is true
                # hence final output is a dictionary with 6 keys (each sector) and each value is a list of finefundamental type objects
                top = sorted(group_by_code[code], key=lambda x: x.DollarVolume, reverse = True)[:ceil(len(group_by_code[code])*percent)]
                # sort these finefundamental type objects by the property of DollarVolume we created earlier and cut it down by the appropriate perecent calculated earlier
                top_list.append(top) # append these stocks to the top_list, where each element is a FineFundamentalType object
                # top_list is a list of 6 lists of finefundamental type objects
            joined_list = top_list[0] # this will be the first element of the list, which is a list of finefundamental type objects
            for ls in top_list[1:]: # take the 2nd element onwards of the top_list, and iterate through it
                joined_list  += ls # add each of these back to the joined_list. This changes a list of 6 lists, to a list of about 500 finefundamental type objects
            self.symbols = [x.Symbol for x in joined_list][:self.num_fine] # get the symbol object for each element in joined_list (they are finefundamentaltype objects) 
            self.Log(",".join(sorted(i.Value for i in self.symbols))) # uses logging to display

            # reset the flags
            self.rebalance_flag = 0
            self.first_month_trade_flag = 0
            return self.symbols # return the symbol list
        else:
            return self.symbols


    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.

        Arguments:
            data: Slice object keyed by symbol containing the stock data
        '''
        pass
    
    def monthly_rebalance(self): # at midnight at start of month, turn rebalance flag on which will make sure that universe selection is carried out
        self.rebalance_flag = 1
        
    def rebalance(self): # will carry out the actual rebalancing
        
        # We will first do it without a market gate, add the market gate in later
        
        if self.symbols is None: return # if nothing in symbol list, exit the function loop, no point rebalancing
    
        chosen_df = self.calc_return(self.symbols) # this is a function we have defined below that returns a pandas dataframe/series with returns as values and stocks as index
        chosen_df = chosen_df.iloc[:self.num_stocks] # slice the top 30 stocks only
        
        # initialize
        self.existing_pos = 0 
        add_symbols = []
        
        for symbol in self.Portfolio.Keys: # each key is a symbol object - unique security identifier (not a string)
            # not going to include code for SPY, as we should never have SPY
            if (symbol.Value not in chosen_df.index): # if we are holding a ticker not in the chosen list (both as strings)
                self.SetHoldings(symbol, 0) # exit trade, note we give .SetHoldings() method the symbol object, not a string
            elif (symbol.Value in chosen_df.index): # if we are golding a ticker that is in the chosen list (both as strings)
                self.existing_pos += 1 # increment the counter, but we haven't rebalanced yet
        
        # calculate the weight that each position should hold
        weight = 0.99/len(chosen_df)
        
        for symbol in chosen_df.index: # loop through symbol objects in chosen list, symbol is now a string
            self.AddEquity(symbol) # add equity before trading it, note that we are giving it a string, but that is ok here, it will add the security identifier that is current for that string
            self.SetHoldings(symbol, weight) # adjust holdings (either existing or from 0) to the appropraite weight.
            
        
    # STILL NEED TO WRITE A FUNCTION TO CALCULATE MODIFIED MOMENTUM AS WELL.
    
    def calc_return(self, stocks): # give this a list QC symbol class objects (essentially the QC500 that was generated in coarse and fine fundamental filter
    # and calculate the traditional momentum 
        # Need to generate a price history and current price
        hist = self.History(stocks, 200, Resolution.Daily) # get 200 day data History
        current = self.History(stocks, 1, Resolution.Minute) # get last price
        
        self.price = {}
        ret = {}
        
        for symbol in stocks: # each symbol in the QC500
            if str(symbol) in hist.index.levels[0] and str(symbol) in current.index.levels[0]: # level 0 = symbol, level 1 = date/time, columns OHLCV, etc
                # so if we have available history and current data for the symbol
                # append the price data to the dictionary self.price
                # by adding a key/value pair where key is symbol as a string, value is the pandas series of closing prices for last 200 days, converted to a list
                self.price[symbol.Value] = list(hist.loc[str(symbol)]['close'])
                # to that key/value pair, append the current closing price
                # need to use [0] even though it is a point because it is a pandas series with 1 element, using [0] converts that to a float
                self.price[symbol.Value].append(current.loc[str(symbol)]['close'][0]) # append the current value to the dictionary
                
        for symbol in self.price.keys(): # loop through dictionary keys (stock names as strings) and call that symbol
            ret[symbol] = (self.price[symbol][-1] - self.price[symbol][0]) / self.price[symbol][0]
            # calculate the returns over last 200 days as a percentage, last price less first price / first price
            # create a dictionary where key = symbol as a string, value = return as a float
            
        df_ret = pd.DataFrame.from_dict(ret, orient='index')
        # make a dataframe from the dictionary ret with symbol (string) / returns (float) as key/value
        # by specifying orient = 'index' we make the keys the index and the returns will be the column
        df_ret.columns = ['return'] # name the column
        sort_return = df_ret.sort_values(by = ['return'], ascending = False) # sort the returns highest to lowest
        
        return sort_return # returns a dataframe with the returns in order of highest to lowest
            

        # NEXT STEP IS TO COMPLETE A TRADITIONAL MOMENTUM ALGORITHM, RUN IT AND BACKTEST 
        # THEN ADD A MARKET GATE 
        # THEN ADD A MODIFIED MOMENTUM CALCULATION