import pandas as pd
from datetime import datetime

class AssetClassTrendFollowingAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        self.SetStartDate(2008, 5, 1)    # Set Start Date
        self.SetEndDate(datetime.now())  # Set End Date
        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 ["SPY", "EFA", "BND", "VNQ", "GSG"]]

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

        self.SetAlpha(AssetClassTrendFollowingAlphaModel(self))

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) 
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(NullRiskManagementModel())
    
class AssetClassTrendFollowingAlphaModel(AlphaModel):
    '''Alpha model that uses an SMA and security price to create insights'''

    def __init__(self,
                 algorithm,
                 period = 10*21,
                 resolution = Resolution.Daily):
        '''
        Initializes a new instance of the AssetClassTrendFollowingAlphaModel class
        Args:
            algorithm: The algorithm instance
            period: The SMA period
            resolution: The data resolution
        '''
        self.period = period
        self.resolution = resolution
        self.smaBySymbol = dict()

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = f'{self.__class__.__name__}({period},{resolutionString})'

        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), period)


    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
        '''
        
        insights = []

        for symbol,sma in self.smaBySymbol.items():
            if sma.IsReady:
                if algorithm.Securities[symbol].Price > sma.Current.Value:
                    insights.append(Insight.Price(symbol,self.predictionInterval,InsightDirection.Up))

        return insights
            

    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.smaBySymbol.setdefault(symbol, algorithm.SMA(symbol, self.period, self.resolution))
            for time, value in history[str(symbol)].iteritems():
                data.Update(time, value)