import pandas as pd

class MomentumAndStyleRotationAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        self.SetStartDate(2018, 1, 1) 
        self.SetEndDate(2019, 3, 1)    
        self.SetCash(100000)             # Set Strategy Cash
        
        self.SetSecurityInitializer(lambda security: security.SetFeeModel(ConstantFeeModel(0)))

        symbols = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) 
            for ticker in [ "IJJ", # iShares S&P MidCap 400 Value Index ETF
                            "IJS", # iShares S&P SmallCap 600 Value ETF 
                            "IVE", # iShares S&P 500 Value Index ETF 
                            "IVW", # iShares S&P 500 Growth ETF   
                            "IJK", # iShares S&P Mid-Cap 400 Growth ETF
                            "IJT", # iShares S&P Small-Cap 600 Growth ETF
                          ]]

        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )

        self.SetAlpha(MomentumAndStyleRotationAlphaModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) 
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())
    
class MomentumAndStyleRotationAlphaModel(AlphaModel):
    '''Alpha model that uses the Momentum Indicator to create insights'''

    def __init__(self,
                 period = 12*20,
                 resolution = Resolution.Daily):
        '''
        Initializes a new instance of the MomentumAndStyleRotationAlphaModel class
        Args:
            period: The Momentum lookback period
            resolution: The data resolution
        '''
        self.period = period
        self.resolution = resolution
        self.MomentumBySymbol = dict()

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = f'{self.__class__.__name__}({period},{resolutionString})'
        
        self.lastMonth = -1
        
    def Update(self, algorithm, data):
        '''
        Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated
        '''
        month = algorithm.Time.month
        if month == self.lastMonth:
            return []
        self.lastMonth = month

        momentumBySymbol = {symbol: value for symbol, value in self.MomentumBySymbol.items() if value.IsReady}
        sortedMomentum = [x[0] for x in sorted(momentumBySymbol.items(), key = lambda kv: kv[1], reverse=True)]
        
        expiry = algorithm.Time.replace(month = month + 1, day = 1) if month < 12 else \
            algorithm.Time.replace(year = algorithm.Time.year + 1, month = 1, day = 1)
            
        return Insight.Group(
            [
                Insight.Price(sortedMomentum[0], expiry, InsightDirection.Up),
                Insight.Price(sortedMomentum[-1], expiry, InsightDirection.Down)
            ])

    def OnSecuritiesChanged(self, algorithm, changes):
        '''
        Event fired each time the we add securities from the data feed
        
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm
        '''
        addedSymbols = [ x.Symbol for x in changes.AddedSecurities ]
        history = algorithm.History(addedSymbols, self.period, self.resolution)
        history = history.close.unstack(level=0)
        
        for symbol in addedSymbols:
            data = self.MomentumBySymbol.setdefault(symbol, algorithm.MOM(symbol, self.period, self.resolution))
            for time, value in history[str(symbol)].iteritems():
                data.Update(time, value)