from dateutil.relativedelta import relativedelta

class FF(PythonData):
    """
    This class is used to stream Fama French data into our algorithm.
    """
    
    def GetSource(self, config, date, isLiveMode):
        """
        Return the URL string source of the file. This will be converted to a stream 
        
        Inputs:
         - config
            Configuration object
         - date
            Date of this source file
         - isLiveMode
            True if we're in live mode; False for backtesting mode
        
        Returns a SubscriptionDataSource - the source location and transport medium for a subscription.
        """
        source = "https://github.com/QuantConnect/Tutorials/raw/feature-data-directory/Data/F-F_Research_Data_Factors.csv"
        return SubscriptionDataSource(source, SubscriptionTransportMedium.RemoteFile)
    
    def Reader(self, config, line, date, isLive):
        """
        Reader converts each line of the data source into BaseData objects. Each data type creates its own 
        factory method, and returns a new instance of the object each time it is called. The returned object 
        is assumed to be time stamped in the config.ExchangeTimeZone.
        
        Inputs:
         - config
            Subscription data config setup object
         - line
            Line of the source document
         - date
            Date of the requested data
         - isLive
            True if we're in live mode; False for backtesting mode
        
        Returns a data point from the Fama French data feed.
        """
        # If first character is not digit, pass
        if not (line.strip() and line[0].isdigit()): 
            return None
    
        try:
            data = line.split(',')

            ff = FF()
            ff.Symbol = config.Symbol
            ff.Time = datetime.strptime(data[0], '%Y%m') + relativedelta(months=1)

            ff.SetProperty("hml", float(data[3]))
            ff.SetProperty("mkt", float(data[1]))
            ff.SetProperty("rf", float(data[4]))
            ff.SetProperty("smb", float(data[2]))
            
            return ff
        except ValueError:
            # Do nothing, possible error in json decoding
            return None

import statsmodels.api as sm
import pandas as pd

class ResidualMomentum:
    """
    This class manages a rolling window of the previous `num_train_months` months. It gathers
    its data via a consolidator into monthly bars. Every month, a regression model is fit to the 
    residual returns over the previous `num_train_months` months (t-`num_train_months` - t-1). 
    It calculates a score based on the residual returns over the previous `num_test_months` 
    months (excluding the lastest month) (t-`num_test_months` - t-2). 
    """
    
    def __init__(self, symbol, monthly_returns, algorithm, alpha, close, num_train_months, num_test_months, min_price):
        """
        Inputs:
         - symbol
            The symbol to apply the indicator on
         - monthly_returns
            Trailing monthly returns for the symbol
         - algorithm
            Algorithm instance running the backtest
         - alpha
            Refrence to the ResidualMomentumAlphaModel.
         - close
            Closing price of the latest full month
         - num_train_months
            Number of months to train the regression model (> 2)
         - num_test_months
            Number of months to test the regression model (1 < num_test_months < num_train_months)
         - min_price
            Minimum price a security needs to be considered in the rebalance (>= 0)
        """
        self.Symbol = symbol
        self.monthly_returns = monthly_returns
        self.monthly_returns.columns = ['m_return']
        self.score = None
        self.alpha = alpha
        self.num_train_months = num_train_months
        self.num_test_months = num_test_months
        self.min_price = min_price
        
        # Setup monthly consolidation
        self.consolidator = TradeBarConsolidator(self.CustomMonthly)
        self.consolidator.DataConsolidated += self.CustomMonthlyHandler
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
        # Set the initial score
        self.update_score(close)
    
    def dispose(self, algorithm):
        """
        Removes the monthly conoslidator.
        
        Inputs
         - algorithm
            The QCAlgorithm object
        """
        algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator)
    
    def update_score(self, close):
        """
        Updates the score for the Residual Momentum indicator.
        
        Inputs
         - close
            The closing price of the latest full month
        """
        # If current price < $`min_price`, don't bother calculating the score
        if close < self.min_price:
            self.score = None
            return
        
        # Fit regression model over the previous `num_train_months` months
        X = self.alpha.fama_french_factors
        X = sm.add_constant(X)
        y = self.monthly_returns.values
        model = sm.OLS(y, X).fit()

        # Calculate score on the previous `num_test_months` months, excluding the most recent month
        # (t-`num_test_months` - t-2)
        pred = model.predict(X.iloc[-self.num_test_months:-1])
        residual_returns = self.monthly_returns.iloc[-self.num_test_months:-1].values - pred.values
        self.score = residual_returns.sum() / residual_returns.std()

    def CustomMonthly(self, dt):
        '''Custom Monthly Func'''
        start = dt.replace(day=1).date()
        end = dt.replace(day=28) + timedelta(4)
        end = (end - timedelta(end.day-1)).date()
        return CalendarInfo(start, end - start)

    def CustomMonthlyHandler(self, sender, consolidated):
        """
        Updates the monthly returns rolling window and the score.
        
        Inputs
         - sender
            Function calling the consolidator
         - consolidated
            Tradebar representing the latest completed month
        """
        # Append to monthly returns rolling window DataFrame
        monthly_return = (consolidated.Close - consolidated.Open) / consolidated.Close
        row = pd.DataFrame({'m_return' : [monthly_return]}, index=[consolidated.Time])
        self.monthly_returns = self.monthly_returns.append(row).iloc[-self.num_train_months:]

        self.update_score(consolidated.Close)

import pandas as pd
import numpy as np
from dateutil.relativedelta import relativedelta

from FamaFrench import FF
from ResidualMomentum import ResidualMomentum

class ResidualMomentumAlphaModel(AlphaModel):
    """
    This class houses the Fama French data and a dictionary of ResidualMomentum indicators for 
    symbols. Each month, we rank the symbols by the ResidualMomentum indicator scores, then
    emit insights to generate a long-short portfolio with the symbols having the highest and lowest 
    scores.
    """
    symbol_data = {}
    fama_french_factors = pd.DataFrame()
    month = -1

    def __init__(self, algorithm, num_train_months=36, num_test_months=12, long_short_pct=10, min_price=1):
        """
        Inputs:
         - algorithm
            Algorithm instance running the backtest
         - num_train_months
            Number of months to train the regression model (> 2)
         - num_test_months
            Number of months to test the regression model (1 < num_test_months < num_train_months)
         - long_short_pct
            The percentage of the universe we go long and short (0 < long_short_pct <= 50)
         - min_price
            Minimum price a security needs to be considered in the rebalance (>= 0)
        """
        self.num_train_months = num_train_months
        self.num_test_months = num_test_months
        self.long_short_pct = long_short_pct
        self.min_price = min_price
        
        self.ff = algorithm.AddData(FF, "FF", Resolution.Daily).Symbol
        
        # Warmup FF history
        end = algorithm.StartDate - timedelta(1)
        start = end - relativedelta(months=self.num_train_months)
        self.fama_french_factors = algorithm.History(self.ff, start, end).loc[self.ff]

    def Update(self, algorithm, slice):
        """
        Called each time our alpha model receives a new data slice.
        
        Inputs:
         - algorithm
            Algorithm instance running the backtest
         - slice
            A data structure for all of an algorithm's data at a single time step
        
        Returns an empty list or an Insight group to the portfolio construction model
        """
        # If we have a new month of fama french data, update our df
        if slice.ContainsKey(self.ff):
            self.update_ff(slice[self.ff])
        
        if algorithm.IsWarmingUp:
            return []
            
        # Only update insights at the start of every month
        if algorithm.Time.month == self.month:
            return []
        self.month = algorithm.Time.month
        
        # Sort self.symbol_data values by their residual momentum score
        has_score = [s for s in self.symbol_data.values() if s.score is not None]
        sorted_by_score = sorted(has_score, key=lambda x: x.score, reverse=True)
        
        # If the universe is too small to grab `long_short_pct`% on both sides, do nothing.
        num_passed = int(len(sorted_by_score) * (1 / self.long_short_pct))
        if num_passed == 0:
            return []
            
        # Create insights
        insights = []
        # Long the top `long_short_pct`% of symbols, based on score
        for s in sorted_by_score[:num_passed]:
            if s.Symbol in slice.Bars:
                insights.append(Insight(s.Symbol, timedelta(days=30), InsightType.Price, InsightDirection.Up))
        
        # Short the bottom `long_short_pct`% 
        for s in sorted_by_score[-num_passed:]:
            if s.Symbol in slice.Bars:
                insights.append(Insight(s.Symbol, timedelta(days=30), InsightType.Price, InsightDirection.Down))
        
        return Insight.Group(insights)

    def OnSecuritiesChanged(self, algorithm, changes):
        """
        Called each time our universe has changed.
        
        Inputs:
         - algorithm
            Algorithm instance running the backtest
         - changes
            The additions and subtractions to the algorithm's security subscriptions
        """
        if len(changes.AddedSecurities) > 0:
            # Get lookback dates
            end_lookback = Expiry.EndOfMonth(algorithm.Time) - relativedelta(months=1)
            start_lookback = end_lookback - relativedelta(months=self.num_train_months)
            
            # Get history of symbols over lookback window
            added_symbols = [x.Symbol for x in changes.AddedSecurities]
            history = algorithm.History(added_symbols, start_lookback, end_lookback, Resolution.Daily)
            
            # Filter for sufficient history
            if history.shape[0] > 0 and \
                (history.index.levels[1][-1] - history.index.levels[1][0]).days / 30 >= self.num_train_months:
                
                # Get monthly returns of symbols with sufficient history
                monthly_returns, closes = self.calc_performance(changes.AddedSecurities, history)
                
                for added in monthly_returns.columns:
                    # Create residual momentum indicator for this symbol
                    ret = monthly_returns[[added]].iloc[-self.num_train_months:]
                    self.symbol_data[added] = ResidualMomentum(added, ret, algorithm, self, closes[added], 
                                                                self.num_train_months, self.num_test_months, self.min_price)
            
        for removed in changes.RemovedSecurities:
            # Remove symbol from our symbol_data dictionary
            resid_mom = self.symbol_data.pop(removed.Symbol, None)
            if resid_mom:
                # Remove consolidator
                resid_mom.dispose(algorithm)
        
    def calc_performance(self, securities, history):
        """
        Calculates the monthly returns for securites over a historical period.
        Securities with insufficient history are omitted.
        
        Inputs:
        # - securities
        #    List of security objects to calculate the monthly returns for
         - history
            DataFrame containing the historical prices of securities
            
        Returns a DataFrame containing the monthly returns and the latest month's closing 
        price for securities with sufficient history.
        """
        symbols = [x.Symbol for x in securities]
        
        # Must not have null history
        duration_filter = ~history['close'].unstack(level=0).isnull().any()
        duration_filter = duration_filter[duration_filter].index
        history = history.loc[duration_filter].copy()
        
        # Roll back the timestamp of our history DataFrame by one day
        history = history.unstack(level=0)
        history = history.set_index(history.index.map(lambda x: x - timedelta(days=1))).stack().swaplevel()
        
        # Calculate monthly returns
        returns = {sym : [] for sym in symbols if sym in history.index}
        indicies = []
        for i, sym in enumerate(returns):
            for idx, g in history.loc[sym].groupby(pd.Grouper(freq='M')):
                monthly_ret = (g.iloc[-1].close - g.iloc[0].open) / g.iloc[0].open
                returns[sym].append(monthly_ret)
                if i == 0:
                    indicies.append(idx)
        monthly_returns = pd.DataFrame(returns, index=indicies)
        
        # Save latest closing price for each symbol
        closes = {sym : history.loc[sym].iloc[-1].close for sym in returns}
        
        return monthly_returns, closes
        
        return None, None
        
        
    def update_ff(self, data):
        """
        Updates the fama and french DataFrame with the latest data
        
        Inputs:
         - data
            PythonData object containing the ff data
        """
        row = pd.DataFrame({"hml" : [data.GetProperty('hml')],
                            "mkt" : [data.GetProperty('mkt')],
                            "rf"  : [data.GetProperty('rf')],
                            "smb" : [data.GetProperty('smb')]},
                            index=[data.Time])
        self.fama_french_factors = self.fama_french_factors.append(row).iloc[-self.num_train_months:]

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class TopMarketCapUniverseSelection(FundamentalUniverseSelectionModel):
    """
    This universe selection model refreshes monthly to contain the securities which
    have the largest market capitalizations.
    """
    def __init__(self, coarse_size = 400, fine_pct = 10):
        """
        Inputs:
         - coarse_size
            Number of securities to return from coarse selection
         - fine_pct
            Percentage of securities to return from fine selection. In decreasing order by 
            market cap
        """
        self.month = 0
        self.coarse_size = coarse_size
        self.fine_pct = fine_pct
        super().__init__(True)

    def SelectCoarse(self, algorithm, coarse):
        """
        Coarse universe selection is called each day at midnight.
        
        Inputs:
         - algorithm
            Algorithm instance running the backtest
         - coarse
            List of CoarseFundamental objects
            
        Returns the first `coarse_size` symbols that have fundamental data.
        """
        if self.month == algorithm.Time.month:
            return Universe.Unchanged
        return [ x.Symbol for x in coarse if x.HasFundamentalData ][:self.coarse_size]
        
    def FineSelectionFunction(self, algorithm, fine):
        """
        Fine universe selection is performed each day at midnight after `SelectCoarse`.
        
        Inputs:
         - algorithm
            Algorithm instance running the backtest
         - fine
            List of FineFundamental objects that result from `SelectCoarse` processing
        
        Returns a list of symbols for the `fine_pct`% of securities with the largest
        market capitalization.
        """
        self.month = algorithm.Time.month

        # Select the top `self.fine_pct`%, based on market cap
        sorted_mkt_cap = sorted(fine, key=lambda x: x.MarketCap, reverse=True)
        universe_size = int(len(sorted_mkt_cap) * (1 / self.fine_pct))
        return [ x.Symbol for x in sorted_mkt_cap[:universe_size] ]

from ResidualMomentumAlpha import ResidualMomentumAlphaModel
from TopMarketCapUniverseSelection import TopMarketCapUniverseSelection

class ResidualMomentumAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1)
        self.SetCash(100000)
        
        self.SetUniverseSelection(TopMarketCapUniverseSelection(100))
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.AddAlpha(ResidualMomentumAlphaModel(self))
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())