import pandas as pd
from dateutil.relativedelta import relativedelta

class ROCAndNearness:
    """
    This class manages the historical data for a symbol and calculates the ROC and Nearness factors.
    """
    def __init__(self, symbol, algorithm, roc_lookback_months, nearness_lookback_months):
        """
        Input:
         - symbol
            Symbol to apply this indicator to
         - algorithm
            Algorithm instance running the backtest
         - roc_lookback_months
            Number of trailing months to calculate the rate of change over (> 0)
         - nearness_lookback_months
            Number of trailing months to calculate the nearness factor over (> 0)
        """
        self.symbol = symbol
        self.algorithm = algorithm
        self.roc_lookback_months = roc_lookback_months
        self.nearness_lookback_months = nearness_lookback_months
        self.lookback_months = max(roc_lookback_months, nearness_lookback_months)

        # Warm up history
        self.warm_up_history(symbol, algorithm)

        # Setup indicator consolidator
        self.consolidator = TradeBarConsolidator(timedelta(1))
        self.consolidator.DataConsolidated += self.CustomDailyHandler
        algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.consolidator)
        
    def CustomDailyHandler(self, sender, consolidated):
        """
        Updates the rolling lookback window with the latest data.
        
        Inputs
         - sender
            Function calling the consolidator
         - consolidated
            Tradebar representing the latest completed trading day
        """
        # Add new data point to history
        if consolidated.Time not in self.history.index:
            row = pd.DataFrame({'open': consolidated.Open, 'high': consolidated.High, 'close': consolidated.Close},
                                index=[consolidated.Time])
            self.history = self.history.append(row)

        # Remove expired history
        start_lookback = Expiry.EndOfMonth(self.algorithm.Time) - relativedelta(months=self.lookback_months + 1)
        self.history = self.history[self.history.index >= start_lookback]
        
    def dispose(self):
        """
        Removes the monthly conoslidator.
        """
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)
        
    def warm_up_history(self, symbol, algorithm):
        """
        Warms up the `historical_prices_by_symbol` dictionary with historical data as far back as the 
        earliest lookback window.
        
        Input:
         - symbol
            Symbol that needs history warm up
         - algorithm
            Algorithm instance running the backtest
        """
        # Get the historical data
        start_lookback = Expiry.EndOfMonth(algorithm.Time) - relativedelta(months=self.lookback_months + 1)
        history = algorithm.History([symbol], start_lookback, algorithm.Time, Resolution.Daily)
        
        # Rollback history timestamp by 1 day to ensure accurate monthly ROC values
        if history.shape[0] > 0:
            history = history.unstack(level=0)
            history = history.set_index(history.index.map(lambda x: x - timedelta(days=1))).stack().swaplevel()
        
        self.history = history.loc[symbol, ['open', 'high', 'close']] if symbol in history.index else pd.DataFrame()
    
    @property
    def IsReady(self):
        """
        Boolean to signal if the symbol has sufficient history to fill the ROC lookback window.
        """
        if self.history.shape[0] == 0:
            return False
        
        return self.history.index[0] < Expiry.EndOfMonth(self.algorithm.Time) - relativedelta(months=self.roc_lookback_months)
    
    @property
    def roc(self):
        """
        Calculates the rate of change over the ROC lookback window.
        """
        lookback = self.get_lookback(self.roc_lookback_months)
        start_price = lookback.iloc[0].open
        end_price = lookback.iloc[-1].close
        return (end_price - start_price) / start_price 
        
    @property
    def nearness(self):
        """
        Calculates how close the closing price of the nearness lookback window was to its maximum price.
        """
        lookback = self.get_lookback(self.nearness_lookback_months)
        return lookback.iloc[-1].close / lookback.high.max()
    
    
    def get_lookback(self, num_months):
        """
        Slices the historical data into the trailing `num_months` months.
        
        Input:
         - num_months
            Number of trailing months in the lookback window
        
        Returns DataFrame containing data for the lookback window.
        """
        start_lookback = Expiry.EndOfMonth(self.algorithm.Time) - relativedelta(months=num_months + 1)
        end_lookback = Expiry.EndOfMonth(self.algorithm.Time) - relativedelta(months=1)
        
        return self.history[(self.history.index >= start_lookback) & (self.history.index < end_lookback)]

import pandas as pd
from dateutil.relativedelta import relativedelta

from ROCAndNearness import ROCAndNearness

class ROCAndNearnessAlphaModel(AlphaModel):
    """
    This class ranks securities in the universe by their historical rate of change and nearness to trailing highs.
    """
    symbol_data_by_symbol = {}
    month = -1
    
    def __init__(self, roc_lookback_months=6, nearness_lookback_months=12, holding_months=6, pct_long_short=10):
        """
        Input:
         - roc_lookback_months
            Number of trailing months to calculate the rate of change over (> 0)
         - nearness_lookback_months
            Number of trailing months to calculate the nearness factor over (> 0)
         - holding_months
            Number of months to hold positions (> 0)
         - pct_long_short
            The percentage of the universe we go long and short (0 < pct_long_short <= 50)
        """
        if roc_lookback_months <= 0 or nearness_lookback_months <= 0 or holding_months <= 0:
            algorithm.Error(f"Requirement violated:  roc_lookback_months > 0 and nearness_lookback_months > 0 and holding_months > 0")
            algorithm.Quit()
            return
         
        if pct_long_short <= 0 or pct_long_short > 50:
            algorithm.Error(f"Requirement violated: 0 < pct_long_short <= 50")
            algorithm.Quit()
            return
        
        self.roc_lookback_months = roc_lookback_months
        self.nearness_lookback_months = nearness_lookback_months
        self.holding_months = holding_months
        self.pct_long_short = pct_long_short
    
    def Update(self, algorithm, data):
        """
        Called each time our alpha model receives a new data slice.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - data
            A data structure for all of an algorithm's data at a single time step
        
        Returns a list of Insights to the portfolio construction model
        """
        # Emit insights on a monthly basis
        time = algorithm.Time
        if self.month == time.month:
            return []
        self.month = time.month
        
        return self.generate_insights(self.ranked_symbols, time)
    
    
    def OnSecuritiesChanged(self, algorithm, changes):
        """
        Called each time our universe has changed.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - changes
            The additions and subtractions to the algorithm's security subscriptions
        """
        for added in changes.AddedSecurities:
            roc_and_nearness = ROCAndNearness(added.Symbol, algorithm, self.roc_lookback_months, self.nearness_lookback_months)
            self.symbol_data_by_symbol[added.Symbol] = roc_and_nearness

        for removed in changes.RemovedSecurities:
            symbol_data = self.symbol_data_by_symbol.pop(removed.Symbol, None)
            if symbol_data:
                symbol_data.dispose()
    
    @property
    def ranked_symbols(self):
        """
        Returns a list of Symbols ordered in increasing rank. A higher rank is associated with
        relatively greater rate of change over the roc lookback window and a closer closing price 
        to the max price in the nearness lookback window.
        """
        ranking_df = pd.DataFrame()
        
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if symbol_data.IsReady:
                row = pd.DataFrame({'ROC': symbol_data.roc, 'Nearness': symbol_data.nearness}, index=[symbol])
                ranking_df = ranking_df.append(row)
                
        return ranking_df.rank().sum(axis=1).sort_values().index
        
    
    def generate_insights(self, ranked_symbols, time):
        """
        Generates the insights to construct a balanced long-short portfolio.
        
        Input:
         - ranked_symbols
            List of Symbols ordered in increasing rank
         - time
            Current time in the backtest
            
        Returns list of Insights.
        """
        insights = []
        
        num_long_short = int(len(ranked_symbols) * (self.pct_long_short / 100))
        if num_long_short > 0:
            hold_duration = Expiry.EndOfMonth(time) + relativedelta(months=self.holding_months-1, seconds=-1)
            for symbol in ranked_symbols[-num_long_short:]:
                insights.append(Insight.Price(symbol, hold_duration, InsightDirection.Up))
            for symbol in ranked_symbols[:num_long_short]:
                insights.append(Insight.Price(symbol, hold_duration, InsightDirection.Down))
        
        return insights

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class AssetManagementUniverseSelection(FundamentalUniverseSelectionModel):
    """
    This universe selection model refreshes monthly to contain US securities in the asset management industry.
    """
    def __init__(self, coarse_size = 400, fine_size = 100):
        """
        Input:
         - coarse_size
            Number of securities to return from coarse selection
         - fine_size
            Number of securities to return from fine selection
        """
        self.month = -1
        self.coarse_size = coarse_size
        self.fine_size = fine_size
        super().__init__(True)

    def SelectCoarse(self, algorithm, coarse):
        """
        Coarse universe selection is called each day at midnight.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - coarse
            List of CoarseFundamental objects
            
        Returns the first `coarse_size` symbols that have fundamental data and traded in the US market.
        """
        if self.month == algorithm.Time.month:
            return Universe.Unchanged
            
        filtered = [x for x in coarse if (x.HasFundamentalData) and (x.Market == "usa")]
        return [ x.Symbol for x in filtered[:self.coarse_size] ]
        
    def SelectFine(self, algorithm, fine):
        """
        Fine universe selection is performed each day at midnight after `SelectCoarse`.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - fine
            List of FineFundamental objects that result from `SelectCoarse` processing
        
        Returns a list of symbols that passed coarse universe selection and are in the asset management industry.
        """
        self.month = algorithm.Time.month
        
        filtered = [f for f in fine if f.AssetClassification.MorningstarIndustryCode == MorningstarIndustryCode.AssetManagement]
        return [ x.Symbol for x in filtered[:self.fine_size] ]

class NetDirectionWeightedPortfolioConstructionModel(PortfolioConstructionModel):
    """
    This PCM allocates its portfolio based on the net direction of insights for all the symbols. A symbol
    that has two active insights with an up direction will have twice the allocation than a symbol with
    only one. Additionally, a symbol that has an up active insight and a down active insight will have no
    position.
    """
    insights = []
    month = -1
    
    def CreateTargets(self, algorithm, insights):
        """
        Called each time the alpha model emits a list of insights.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - insights
            List of insights
            
        Returns a list of portfolio targets.
        """
        for i in insights:
            self.insights.append(i)
        
        if self.month == algorithm.Time.month:
            return []
        self.month = algorithm.Time.month

        # Classify insights
        active_insights = []
        expired_insights = []
        while (len(self.insights) > 0):
            insight = self.insights.pop()
            (active_insights if insight.IsActive(algorithm.UtcTime) else expired_insights).append(insight)
        self.insights = active_insights

        # Liquidate symbols that have expired insights with no active insights
        active_symbols = set([i.Symbol for i in active_insights])
        expired_symbols = set([i.Symbol for i in expired_insights])
        liquidate_symbols = expired_symbols.difference(active_symbols)
        portfolio_targets = [PortfolioTarget.Percent(algorithm, symbol, 0) for symbol in liquidate_symbols]
        
        # Get net direction by symbol and total number of directional insights
        net_direction_by_symbol, num_directional_insights = self.get_net_direction(active_insights)
        
        # Create portfolio targets for active symbols
        for symbol, net_direction in net_direction_by_symbol.items():
            percent = 0 if num_directional_insights == 0 else net_direction / num_directional_insights
            portfolio_targets.append( PortfolioTarget.Percent(algorithm, symbol, percent) )
        
        return portfolio_targets
        
        
    def get_net_direction(self, insights):
        """
        Determines the net direction of each symbol and the number of active directional insights.
        
        Input:
         - insights
            A list of active insights
        
        Returns a dictionary showing the net direction of each symbol, and the number of directional insights.
        """
        net_direction_by_symbol = {}
        num_directional_insights = 0
        
        for insight in insights:
            symbol = insight.Symbol
            direction = insight.Direction
            if symbol in net_direction_by_symbol:
                net_direction_by_symbol[symbol] += direction
            else:
                net_direction_by_symbol[symbol] = direction
                
            num_directional_insights += abs(direction)
        
        return net_direction_by_symbol, num_directional_insights

from AssetManagementUniverseSelection import AssetManagementUniverseSelection
from ROCAndNearnessAlphaModel import ROCAndNearnessAlphaModel
from NetDirectionWeightedPortfolioConstructionModel import NetDirectionWeightedPortfolioConstructionModel

class AssetManagementFirmMomentum(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 4)
        self.SetEndDate(2020, 6, 1)
        
        self.SetCash(100000)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetUniverseSelection(AssetManagementUniverseSelection())

        self.AddAlpha(ROCAndNearnessAlphaModel())
        
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.SetPortfolioConstruction(NetDirectionWeightedPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())