## A simple m odification to add leverage factor to the InsightWeightingPortfolioConstructionModel
class LeveragePCM(InsightWeightingPortfolioConstructionModel):
    
    leverage = 0.0
    
    def CreateTargets(self, algorithm, insights):
        
        targets = super().CreateTargets(algorithm, insights)
        
        return [PortfolioTarget(x.Symbol, x.Quantity*(1+self.leverage)) for x in targets]

''' An ensemble approach to GEM - Global Equities Momentum.
'''
from alpha_model import GEMEnsembleAlphaModel
from pcm import LeveragePCM

class GlobalTacticalAssetAllocation(QCAlgorithm):
    
    def Initialize(self):
         
        self.SetStartDate(2005, 1, 1)
        #self.SetEndDate(2020, 5, 20)
        self.SetCash(100000) 
        self.Settings.FreePortfolioValuePercentage = 0.02
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        # PNQI, TLT
        tickers = ['SPY', 'EFA', 'TLT'] #for plotting
        us_equity = Symbol.Create('SPY', SecurityType.Equity, Market.USA)
        foreign_equity = Symbol.Create('EFA', SecurityType.Equity, Market.USA)
        bond = Symbol.Create('TLT', SecurityType.Equity, Market.USA)
        symbols = [us_equity, foreign_equity, bond]
   
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection( ManualUniverseSelectionModel(symbols) )

        self.AddAlpha( GEMEnsembleAlphaModel(us_equity, foreign_equity, bond) )
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.SetPortfolioConstruction(LeveragePCM(self.RebalanceFunction, PortfolioBias.Long))
        self.lastRebalanceTime = None
        
        self.SetExecution( ImmediateExecutionModel() ) 
        self.AddRiskManagement( NullRiskManagementModel() )
        
        # Initialise plot
        assetWeightsPlot = Chart('AssetWeights %')
        for ticker in tickers:
            assetWeightsPlot.AddSeries(Series(ticker, SeriesType.Line, f'{ticker}%')) 

    def RebalanceFunction(self, time):

        return Expiry.EndOfMonth(self.Time)
    
    def OnData(self, data):
        # Update Plot
        for kvp in self.Portfolio:
                symbol = kvp.Key
                holding = kvp.Value 
                self.Plot('AssetWeights %', f"{str(holding.Symbol)}%", holding.HoldingsValue/self.Portfolio.TotalPortfolioValue)

import numpy as np

class GEMEnsembleAlphaModel(AlphaModel):
    """ If the S&P 500 had positive returns over the past X-months (positive trend) the strategy allocates to stocks 
    the next month; otherwise it allocates to bonds. 
    When the trend is positive for stocks the strategy holds the equity index with the strongest total return 
    over the same horizon. The Ensemble approach takes the average of all signals.
    """

    def __init__(self, us_equity, foreign_equity, bond, resolution=Resolution.Daily):
        '''Initializes a new instance of the SmaAlphaModel class
        Args:
            resolution: The reolution for our indicators
        '''
        self.us_equity = us_equity
        self.foreign_equity = foreign_equity
        self.bond = bond
        self.resolution = resolution
        self.symbolDataBySymbol = {}
        self.month = -1

    def Update(self, algorithm, data):
        '''This is called each time the algorithm receives data for (@resolution of) subscribed securities
        Returns: The new insights generated.
        THIS: analysis only occurs at month start, so any signals intra-month are disregarded.'''
        
        if self.month == algorithm.Time.month:
            return []
        self.month = algorithm.Time.month

        insights = []
        strategies = {}
        weights = dict.fromkeys([self.us_equity, self.foreign_equity, self.bond], 0)
        for symbol, symbolData in self.symbolDataBySymbol.items():
            strategies[symbol] = np.array([])
            for lookback in symbolData.momp.keys():
                strategies[symbol] = np.append(strategies[symbol], symbolData.momp[lookback].Current.Value)
        
        for value in strategies[self.us_equity]:
            if value >= 0:
                us_wins = sum(strategies[self.us_equity] >= strategies[self.foreign_equity])
                foreign_wins = sum(strategies[self.foreign_equity] > strategies[self.us_equity])
                weights[self.us_equity] += us_wins
                weights[self.foreign_equity] += foreign_wins
            else:
                bonds_win = len(strategies[self.us_equity])
                weights[self.bond] += bonds_win


        # GEM Rules. Go to Bonds if SPY Momentum is -ve.
        #bonds_weight = (strategies[self.us_equity] < 0).sum()  
        # BUY SPY if it's +ve, AND its relative momentum is greater than that of Foreign Equities
        #us_equity_weight = ( (strategies[self.us_equity] >= 0) & (strategies[self.us_equity] >= strategies[self.foreign_equity]) ).sum()
        # Else buy Foreign Equities.
        #foreign_equity_weight = ( (strategies[self.us_equity] > 0) & (strategies[self.foreign_equity] > strategies[self.us_equity]) ).sum()
        
        
        insights.append(Insight.Price(self.us_equity, Expiry.EndOfMonth, InsightDirection.Up, None, None, None, weights[self.us_equity]))
        insights.append(Insight.Price(self.foreign_equity, Expiry.EndOfMonth, InsightDirection.Up, None, None, None, weights[self.foreign_equity]))
        insights.append(Insight.Price(self.bond, Expiry.EndOfMonth, InsightDirection.Up, None, None, None, weights[self.bond]))
            
            
        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        
        for added in changes.AddedSecurities:
            self.symbolDataBySymbol[added.Symbol] = SymbolData(added.Symbol, algorithm, self.resolution)

        for removed in changes.RemovedSecurities:
            symbolData = self.symbolDataBySymbol.pop(removed.Symbol, None)
            if symbolData:
                # Remove consolidator
                symbolData.dispose()          


class SymbolData:
    
    def __init__(self, symbol, algorithm, resolution):
        self.algorithm = algorithm
        self.Symbol = symbol
        self.momp = {}
        for period in range(1, 13):
            self.momp[period] = MomentumPercent(period*21)

        # Warm up Indicators
        history = algorithm.History([self.Symbol], 21*13, resolution).loc[self.Symbol]
        # Use history to build our SMA
        for time, row in history.iterrows():
            for period, momp in self.momp.items():
                self.momp[period].Update(time, row["close"])

        # 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):

        for period, momp in self.momp.items():
            self.momp[period].Update(consolidated.Time, consolidated.Close)

        
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator)