#region imports
from AlgorithmImports import *

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

from fama_french import FF
from residual_momentum import ResidualMomentum
#endregion


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.
    """
    
    fama_french_factors = pd.DataFrame()
    _symbol_data = {}
    _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.add_data(FF, "FF", Resolution.DAILY).symbol
        
        # Warmup FF history
        end = algorithm.start_date - timedelta(1)
        start = end - relativedelta(months=self._num_train_months)
        self._factor_names = ['hml', 'mkt', 'smb']
        self.fama_french_factors = algorithm.history(self._ff, start, end).loc[self._ff][self._factor_names]

    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.contains_key(self._ff):
            self._update_ff(slice[self._ff])
        
        if algorithm.is_warming_up:
            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 on_securities_changed(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.added_securities) > 0:
            # Get lookback dates
            end_lookback = Expiry.end_of_month(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.added_securities]
            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.added_securities, 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.removed_securities:
            # 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
        """
        self.fama_french_factors.loc[data.time, :] = [data.get_property(factor) for factor in self._factor_names]
        self.fama_french_factors = self.fama_french_factors.iloc[-self._num_train_months:]
        

#region imports
from AlgorithmImports import *

from dateutil.relativedelta import relativedelta
#endregion


class FF(PythonData):
    """
    This class is used to stream Fama French data into our algorithm.
    """
    
    def get_source(self, config, date, is_live_mode):
        """
        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.REMOTE_FILE)
    
    def reader(self, config, line, date, is_live):
        """
        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.set_property("hml", float(data[3]))
            ff.set_property("mkt", float(data[1]))
            ff.set_property("rf", float(data[4]))
            ff.set_property("smb", float(data[2]))
            
            return ff
        except ValueError:
            # Do nothing, possible error in json decoding
            return None

#region imports
from AlgorithmImports import *

from alpha import ResidualMomentumAlphaModel
from universe import TopMarketCapUniverseSelection
#endregion


class ResidualMomentumAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2016, 1, 1)
        self.set_end_date(2020, 4, 1)
        self.set_cash(100000)
        
        self.set_universe_selection(TopMarketCapUniverseSelection(100))
        self.universe_settings.resolution = Resolution.DAILY
        
        self.add_alpha(ResidualMomentumAlphaModel(self))
        
        self.set_portfolio_construction(EqualWeightingPortfolioConstructionModel())
        
        self.set_execution(ImmediateExecutionModel())
        

#region imports
from AlgorithmImports import *

import statsmodels.api as sm
import pandas as pd
#endregion


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.score = None
        self.symbol = symbol
        self._monthly_returns = monthly_returns
        self._monthly_returns.columns = ['m_return']
        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._custom_monthly)
        self._consolidator.data_consolidated += self._custom_monthly_handler
        algorithm.subscription_manager.add_consolidator(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.subscription_manager.remove_consolidator(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 _custom_monthly(self, dt):
        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 _custom_monthly_handler(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
        self._monthly_returns.loc[consolidated.time, 'm_return'] = monthly_return
        self._monthly_returns = self._monthly_returns.iloc[-self._num_train_months:]

        self._update_score(consolidated.close)
        

#region imports
from AlgorithmImports import *
#endregion
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 select_coarse(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

        sorted_by_dollar_vol = sorted([x for x in coarse if x.has_fundamental_data],
                                    key=lambda x: x.dollar_volume, reverse=True)
        return [x.symbol for x in sorted_by_dollar_vol][:self._coarse_size]
        
    def fine_selection_function(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.market_cap, reverse=True)
        universe_size = int(len(sorted_mkt_cap) * (1 / self._fine_pct))
        return [ x.symbol for x in sorted_mkt_cap[:universe_size] ]