import pandas as pd

class ETFSharpeRotationAlgorithm(QCAlgorithm):
    # This strategy was sourced from https://www.rotationinvest.com/blog/etf-rotation-strategy-using-sharpe-ratio
    #
    # Rotating Funds:  SPY, EFA, GLD
    # Cash Filter:     TLT
    # Strategy Rules:  Invest in the top 1 fund based on the trailing 3-month Sharpe Ratio
    #                  If the top fund is below it's 150-day moving average then invest in the bond ETF (TLT)

    def Initialize(self):
        self.SetStartDate(2005, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        trading_days_per_month = 21
        self.sharpe_lookback = 3 * trading_days_per_month
        
        sma_lookback = 150
        lookback = max(self.sharpe_lookback, sma_lookback)
        
        rotation_tickers = ['SPY', 'EFA', 'GLD']
        self.sma_by_rotation_symbol = {}
        self.history_by_symbol = {}
        for ticker in rotation_tickers:
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            self.sma_by_rotation_symbol[symbol] = self.SMA(symbol, sma_lookback, Resolution.Daily)
            
            history = self.History(symbol, lookback, Resolution.Daily).loc[symbol].close
            for time, close in history[-lookback:].iteritems():
                self.sma_by_rotation_symbol[symbol].Update(time, close)
            
            self.history_by_symbol[symbol] = history[-self.sharpe_lookback:]
            
            # Setup a consolidator for sharpe calculation
            consolidator = TradeBarConsolidator(timedelta(1))
            consolidator.DataConsolidated += self.CustomDailyHandler
            self.SubscriptionManager.AddConsolidator(symbol, consolidator)
            
        self.cash_filter_symbol = self.AddEquity("TLT", Resolution.Daily).Symbol
        
        self.month = -1
        
    def CustomDailyHandler(self, sender, consolidated):
        row = pd.Series([consolidated.Close], index=[consolidated.Time])
        appended = self.history_by_symbol[consolidated.Symbol].append(row)
        self.history_by_symbol[consolidated.Symbol] = appended.iloc[-self.sharpe_lookback:]

    def OnData(self, data):

        # Ensure we have data
        for symbol, _ in self.sma_by_rotation_symbol.items():
            if not data.ContainsKey(symbol) or data[symbol] is None:
                return
        if not data.ContainsKey(self.cash_filter_symbol) or data[self.cash_filter_symbol] is None:
            return
        
        # Rebalance monthly
        if data.Time.month == self.month:
            return
        self.month = data.Time.month
        
        # Select the symbol with the highest trailing sharpe ratio
        sharpe_by_symbol = {}
        for symbol, history in self.history_by_symbol.items():
            sharpe = history.pct_change().mean() / history.pct_change().std()
            sharpe_by_symbol[symbol] = sharpe
        top_symbol = sorted(sharpe_by_symbol.items(), key=lambda item: item[1], reverse=True)[0][0]
        
        # If the top_symbol is above it's SMA, invest it in. Otherwise, invest in the cash filter (TLT)
        sma = self.sma_by_rotation_symbol[top_symbol].Current.Value
        if data[symbol].Price >= sma:
            for symbol, _ in self.sma_by_rotation_symbol.items():
                if self.Securities[symbol].Invested and symbol != top_symbol:
                    self.Liquidate(symbol)
            if self.Securities[self.cash_filter_symbol].Invested:
                self.Liquidate(self.cash_filter_symbol)
            
            self.SetHoldings(top_symbol, 1)
        else:
            for symbol, _ in self.sma_by_rotation_symbol.items():
                self.Liquidate(symbol)
            self.SetHoldings(self.cash_filter_symbol, 1)