# https://quantpedia.com/strategies/fed-model/
#
# Each month, the investor conducts a one-month predictive regression (using all available data up to that date) predicting excess stock market
# returns using the yield gap as an independent variable. The “Yield gap” is calculated as YG = EY − y, with earnings yield EY ≡ ln (1 ++ E/P)
# and y = ln (1 ++ Y) is the log 10 year Treasury bond yield. Then, the strategy allocates 100% in the risky asset if the forecasted excess
# returns are positive, and otherwise, it invests 100% in the risk-free rate.

from collections import deque
import numpy as np
from scipy import stats

class FEDModel(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)
        self.SetCash(100000)
        
        # monthly price data and yield gap data
        self.data = {}
        
        self.period = 12 * 21
        self.SetWarmUp(self.period)

        self.market = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.market_data = deque()
        
        self.cash = self.AddEquity('SHY', Resolution.Daily).Symbol
        
        # risk free rate
        self.risk_free_rate = self.AddData(QuandlValue, 'FRED/DGS3MO', Resolution.Daily).Symbol
        
        # 10Y bond yield symbol
        self.bond_yield = self.AddData(QuantpediaBondYield, 'US10YT', Resolution.Daily).Symbol
        
        # SP500 earnings yield data
        self.sp_earnings_yield = self.AddData(QuandlValue, 'MULTPL/SP500_EARNINGS_YIELD_MONTH', Resolution.Daily).Symbol
        
        self.yield_gap = deque()
        
        self.recent_month = -1
    
    def OnData(self, data):
        rebalance_flag = False
        
        if self.sp_earnings_yield in data and data[self.sp_earnings_yield]:
            if self.Time.month != self.recent_month:
                self.recent_month = self.Time.month
                rebalance_flag = True

        if not rebalance_flag:
            # earnings yield data is no longer comming in
            if self.Securities[self.sp_earnings_yield].GetLastData():
                if (self.Time.date() - self.Securities[self.sp_earnings_yield].GetLastData().Time.date()).days > 31:
                    self.Liquidate()
            return
        
        # pdate market price data
        if self.market in data and self.risk_free_rate in data and self.bond_yield in data:
            if data[self.market] and data[self.risk_free_rate] and data[self.bond_yield]:
                market_price = data[self.market].Value
                rf_rate = data[self.risk_free_rate].Value
                bond_yield = data[self.bond_yield].Value
                sp_ey = data[self.sp_earnings_yield].Value
                if market_price != 0 and rf_rate != 0 and bond_yield != 0 and sp_ey != 0:
                    self.market_data.append((market_price, rf_rate))

                    yield_gap = np.log(sp_ey) - np.log(bond_yield)
                    self.yield_gap.append(yield_gap)
                    rebalance_flag = True
        
        
        # ensure minimum data points to calculate regression
        min_count = 6
        if len(self.market_data) >= min_count:
            market_closes = np.array([x[0] for x in self.market_data])
            market_returns = (market_closes[1:] - market_closes[:-1]) / market_closes[:-1]
            rf_rates = np.array([x[1] for x in self.market_data][1:])
            excess_returns = market_returns - rf_rates
            
            yield_gaps = [x for x in self.yield_gap]

            # linear regression
            # Y = α + (β ∗ X)
            # intercept = alpha
            # slope = beta
            beta, alpha, r_value, p_value, std_err = stats.linregress(yield_gaps[1:-1], market_returns[1:])
            X = yield_gaps[-1]
            
            # predicted market return
            Y = alpha + (beta * X)  
            
            # trade execution / rebalance
            if Y > 0:
                if self.Portfolio[self.cash].Invested:
                    self.Liquidate(self.cash)
                self.SetHoldings(self.market, 1)
            else:
                if self.Portfolio[self.market].Invested:
                    self.Liquidate(self.market)
                self.SetHoldings(self.cash, 1)

# Quantpedia bond yield data.
# NOTE: IMPORTANT: Data order must be ascending (datewise)
class QuantpediaBondYield(PythonData):
    def GetSource(self, config, date, isLiveMode):
        return SubscriptionDataSource("data.quantpedia.com/backtesting_data/bond_yield/{0}.csv".format(config.Symbol.Value), SubscriptionTransportMedium.RemoteFile, FileFormat.Csv)

    def Reader(self, config, line, date, isLiveMode):
        data = QuantpediaBondYield()
        data.Symbol = config.Symbol
        
        if not line[0].isdigit(): return None
        split = line.split(',')
        
        data.Time = datetime.strptime(split[0], "%Y-%m-%d") + timedelta(days=1)
        data['yield'] = float(split[1])
        data.Value = float(split[1])

        return data

# Quandl "value" data
class QuandlValue(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = 'Value'