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

import datetime
import pandas as pd

class MomentumAndStyleRotationAlgorithm(QCAlgorithmFramework):

    def Initialize(self):

        self.SetStartDate(2018, 1, 1)    # Set Start Date
        self.SetEndDate(2018, 2, 1)      # Set End Date
        self.SetCash(100000)             # Set Strategy Cash
        
        self.SetSecurityInitializer(lambda security: security.SetFeeModel(ConstantFeeModel(0)))
    
        self.tickers = [
                        "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
        symbols = [ Symbol.Create(ticker, SecurityType.Equity, Market.USA) for ticker in self.tickers ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        
        self.SetAlpha(MomentumAndStyleRotationAlphaModel())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) 
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetRiskManagement(NullRiskManagementModel())
        
class MomentumAndStyleRotationAlphaModel(AlphaModel):
    '''Alpha model that uses an SMA and security price to create insights'''

    def __init__(self,
                 period = 12*20,
                 resolution = Resolution.Daily,
                 datetime = datetime,
                 pd = pd):
        
        '''
        Initializes a new instance of the MomentumAndStyleRotationAlphaModel 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)
        
        self.lastMonth = -1
        self.longSymbol = None
        self.shortSymbol = None
        self.monthStart = False
        
        self.longStockList = []
        self.shortStockList = []


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

        month = algorithm.Time.month
        if month == self.lastMonth:
            return insights
            
        self.lastMonth = algorithm.Time.month

        mom = {}
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.mom.IsReady:
                mom.update({symbol : symbolData.momValue})

        sortedMomentum = list(dict(sorted(mom.items(), key=lambda kv: kv[1],reverse=True)))
        
        self.longSymbol = list(sortedMomentum)[0]
        self.shortSymbol = list(sortedMomentum)[-1]
        
        predictionInterval = Time.Multiply(Extensions.ToTimeSpan(self.resolution), self.daysLeftInMonth(year = algorithm.Time.year,month = algorithm.Time.month,day = algorithm.Time.day))
        
        # EMIT INSIGHTS
        insights.append(Insight.Price(self.longSymbol, predictionInterval, InsightDirection.Up))
        insights.append(Insight.Price(self.shortSymbol, predictionInterval, InsightDirection.Down))
            
        return insights
    
    # Calculates number of days left in month
    def daysLeftInMonth(self,year,month,day):
        date = datetime.date(year,month,day)
        if date.month == 12:
            return abs(datetime.date(year,month,day)-date.replace(day=31)).days
        return abs(datetime.date(year,month,day)-date.replace(month=date.month+1, day=1) - datetime.timedelta(days=1)).days


    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.mom = algorithm.MOM(symbol, period, resolution)
        
    def WarmUpIndicators(self, history):
        for tuple in history.itertuples():
            self.mom.Update(tuple.Index, tuple.close)
        
    @property
    def momValue(self):
        return float(self.mom.Current.Value)