from clr import AddReference
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Indicators")

from QuantConnect import *
from QuantConnect.Indicators import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Alphas import *

from Risk.NullRiskManagementModel import NullRiskManagementModel

from datetime import datetime

class AssetClassTrendFollowingAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        self.SetStartDate(2007, 5, 1)    # Set Start Date
        self.SetEndDate(datetime.now())  # Set End Date
        self.SetCash(100000)             # Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        symbols = [ Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in ["SPY", "EFA", "BND", "VNQ", "GSG"] ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        
        self.SetAlpha(AssetClassTrendFollowingAlphaModel())
        
        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,
                 period = 10*21,
                 resolution = Resolution.Daily):
        '''
        Initializes a new instance of the AssetClassTrendFollowingAlphaModel class
        Args:
            period: The SMA period
            '''
        self.period = period
        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.symbolDataBySymbol = {}

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


    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, symbolData in self.symbolDataBySymbol.items():
            if symbolData.sma.IsReady:
                # Emit "up" insight if price exceeds the simple moving average
                if algorithm.Securities[symbol].Price > symbolData.smaValue:
                    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
        '''
        
        # initialize data for added securities
        addedSymbols = [ x.Symbol for x in changes.AddedSecurities ]
        history = algorithm.History(addedSymbols, self.period, self.resolution)
            
        for added in changes.AddedSecurities:
            algorithm.Log(added)

            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                # Create SymbolData objects
                symbolData = SymbolData(symbol = str(added.Symbol),period = self.period,resolution = self.resolution,algorithm = algorithm)
                
                # Warmup indicators
                ticker = SymbolCache.GetTicker(added.Symbol)
                symbolData.WarmUpIndicators(history.loc[ticker])
                
                # Add object to dictionary
                self.symbolDataBySymbol[added.Symbol] = symbolData

class SymbolData:
    '''
    Contains data specific to a symbol required by this model
    '''
    def __init__(self,symbol, period,resolution,algorithm):
        self.symbol = symbol
        self.algorithm = algorithm
        
        self.sma = self.algorithm.SMA(symbol, period, resolution)
        
    def WarmUpIndicators(self, history):
        for tuple in history.itertuples():
            self.sma.Update(tuple.Index, tuple.close)
        
    @property
    def smaValue(self):
        return float(self.sma.Current.Value)