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.
        """
        return self.get_lookback(self.roc_lookback_months).shape[0] > 1
    
    @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=25):
        """
        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.Quit(f"Requirement violated:  roc_lookback_months > 0 and nearness_lookback_months > 0 and holding_months > 0")
        
        if pct_long_short <= 0 or pct_long_short > 50:
            algorithm.Quit(f"Requirement violated: 0 < pct_long_short <= 50")
        
        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
        
        
        # Rank symbols. A higher rank (or index in the list) 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 data.ContainsKey(symbol) and symbol_data.IsReady:
                row = pd.DataFrame({'ROC': symbol_data.roc, 'Nearness': symbol_data.nearness}, index=[symbol])
                ranking_df = ranking_df.append(row)
        ranked_symbols =  ranking_df.rank().sum(axis=1).sort_values().index
        
        # Generate insights to form a balanced long-short portfolio
        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
    
    
    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()

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):
        self.month = -1
        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 symbols that have fundamental data.
        """
        if self.month == algorithm.Time.month:
            return Universe.Unchanged
            
        return [x.Symbol for x in coarse if x.HasFundamentalData]
        
    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 are in the asset management industry.
        """
        self.month = algorithm.Time.month
        
        return [f.Symbol for f in fine if f.AssetClassification.MorningstarIndustryCode == MorningstarIndustryCode.AssetManagement]

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. This PCM doesn't liquidate securities when they are removed from the universe. If it's 
    removed from the universe, it will remain invested until all the security's insights expire.
    """
    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

        # Remove insights of delisted symbols
        self.insights = [i for i in self.insights if algorithm.Securities[i.Symbol].IsTradable]
        
        # 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

# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *

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

class AssetManagementFirmMomentum(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(2020, 8, 16)
        
        self.SetCash(1000000)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.SetUniverseSelection(AssetManagementUniverseSelection())
        
        self.SetAlpha(ROCAndNearnessAlphaModel())
        
        self.SetPortfolioConstruction(NetDirectionWeightedPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())