from QuantConnect.Indicators import SimpleMovingAverage

class MovingAverageAlphaModel(AlphaModel):
    '''Alpha model based on Price crossing an SMA'''

    def __init__(self, smaLength=200, resolution=Resolution.Daily):
        '''Initializes a new instance of the SmaAlphaModel class
        Args:
            period: The SMA period
            resolution: The reolution for the SMA'''
        self.smaLength = smaLength
        self.resolution = resolution
        self.symbolDataBySymbol = {}
        self.month = -1
        self.riskOffAsset = "IEF"

    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 = []
        riskOffWeight = riskOnWeight = 1 / len(self.symbolDataBySymbol)
        
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbol.Value == self.riskOffAsset:
                continue
            price = algorithm.Securities[symbol].Price

            if price != 0 and symbolData.MovingAverage.IsReady:

                if price > symbolData.MovingAverage.Current.Value:
                    insights.append( Insight.Price(symbol, Expiry.EndOfMonth, InsightDirection.Up,  None, None, None, riskOnWeight))

                elif price < symbolData.MovingAverage.Current.Value:
                    insights.append( Insight.Price(symbol, Expiry.EndOfMonth, InsightDirection.Flat,  None, None, None, 0) )
                    riskOffWeight += riskOnWeight
                    
        insights.append( Insight.Price(self.riskOffAsset, Expiry.EndOfMonth, InsightDirection.Up,  None, None, None, riskOffWeight) )
        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        
        for added in changes.AddedSecurities:
            self.symbolDataBySymbol[added.Symbol] = SymbolData(added, algorithm, self.smaLength, 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, security, algorithm, smaLength, resolution):
        self.Security = security
        self.Symbol = security.Symbol
        self.MovingAverage = SimpleMovingAverage(smaLength)
        self.algorithm = algorithm

        # Warm up MA
        history = algorithm.History([self.Symbol], smaLength, resolution).loc[self.Symbol]
        for time, row in history.iterrows():
            self.MovingAverage.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):
        self.MovingAverage.Update(consolidated.Time, consolidated.Close)
        
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator)

'''An implementation of Meb Faber's base model: Global Tactical Asset Allocation model (GTAA)(5) 
with 10-month SimpleMovingAverage Filter (200day) and (monthly rebalance), as found in the paper: 
https://papers.ssrn.com/sol3/papers.cfm?abstract_id=962461
"A Quantitative Approach to Tactical Asset Allocation" published May 2006.
Analysis only occurs at month End/Start, signals are NOT generated intra-month.
'''
from alpha_model import MovingAverageAlphaModel

class GlobalTacticalAssetAllocation(QCAlgorithm):
    
    def Initialize(self):
        
        self.SetStartDate(2011, 10, 29) #self.SetStartDate(2016, 1, 1) 
        self.SetEndDate(2020, 5, 20)
        self.SetCash(100000) 
        self.Settings.FreePortfolioValuePercentage = 0.02
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        
        self.UniverseSettings.Resolution = Resolution.Daily
        symbols = [Symbol.Create(ticker, SecurityType.Equity, Market.USA) 
                    for ticker in [ "SPY", # US Large Cap ETF
                                    "VEA", # Developed Foreign Stocks (TradedSince: 2007/8)ETF 
                                    "IEF", # US 10Y Gov.Bonds ETF 
                                    "DBC", # GSCI Commodities ETF (TradedSince: 2006/3)
                                    "VNQ"  # US RealEstate ETF
                                    ]]
        self.AddUniverseSelection( ManualUniverseSelectionModel(symbols) )
        
        self.AddAlpha( MovingAverageAlphaModel() )
        self.Settings.RebalancePortfolioOnInsightChanges = False
        self.Settings.RebalancePortfolioOnSecurityChanges = False  
        
        self.SetPortfolioConstruction( InsightWeightingPortfolioConstructionModel(self.DateRules.MonthStart('SPY'), PortfolioBias.Long) )

        self.SetExecution( ImmediateExecutionModel() ) 
        
        self.AddRiskManagement( NullRiskManagementModel() )

class YourAlgorithm(QCAlgorithm):
    
    def Initialize(self):
 
        # 1) Setting a Benchmark to plot with equity
        self.benchmarkTicker = 'SPY'
        self.SetBenchmark(self.benchmarkTicker)
        self.initBenchmarkPrice = None

    
    def UpdateBenchmarkValue(self):
        ''' Simulate buy and hold the Benchmark '''
        if self.initBenchmarkPrice is None:
            self.initBenchmarkCash = self.Portfolio.Cash
            self.initBenchmarkPrice = self.Benchmark.Evaluate(self.Time)
            self.benchmarkValue = self.initBenchmarkCash
        else:
            currentBenchmarkPrice = self.Benchmark.Evaluate(self.Time)
            self.benchmarkValue = (currentBenchmarkPrice / self.initBenchmarkPrice) * self.initBenchmarkCash    
            
            
    def OnData(self, data):
        # 2) simulate buy and hold the benchmark and plot its daily value as we are using daily data.
        # Otherwise schedule when to call this function!
        self.UpdateBenchmarkValue()
        self.Plot('Strategy Equity', self.benchmarkTicker, self.benchmarkValue)
### _________________________________________________        
        # Plotting HoldingValues
        for kvp in self.Portfolio:
            symbol = kvp.Key
            holding = kvp.Value 
            self.Debug(str(holding.Symbol))
            self.Plot(f"{str(holding.Symbol)}HoldingValues", holding.HoldingsValue)
### _________________________________________________