# 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
from AlgorithmImports import *
from typing import List, Tuple, Deque
import numpy as np
from scipy import stats

class FEDModel(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2000, 1, 1)
        self.SetCash(100000)
        
        self.period:int = 12 * 21
        self.SetWarmUp(self.period)

        self.market:Symbol = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.market_data:Deque[Tuple[float,float]] = deque()
        
        self.cash:Symbol = self.AddEquity('SHY', Resolution.Daily).Symbol
        
        # risk free rate
        self.risk_free_rate:Symbol = self.AddData(QuandlValue, 'OECD/KEI_IR3TIB01_USA_ST_M', Resolution.Daily).Symbol
        
        # 10Y bond yield symbol
        self.bond_yield:Symbol = self.AddData(QuantpediaBondYield, 'US10YT', Resolution.Daily).Symbol
        
        # SP500 earnings yield data
        self.sp_earnings_yield:Symbol = self.AddData(QuandlValue, 'MULTPL/SP500_EARNINGS_YIELD_MONTH', Resolution.Daily).Symbol
        
        self.yield_gap:Deque[float] = deque()
        
        self.recent_month:int = -1
    
    def OnData(self, data:Slice) -> None:
        rebalance_flag:bool = 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

        # update market price data
        if self.market in data and self.bond_yield in data:
            if data[self.market] and data[self.bond_yield] and self.Securities[self.risk_free_rate].GetLastData():
                market_price:float = data[self.market].Value
                rf_rate:float = self.Securities[self.risk_free_rate].Price
                bond_yield:float = data[self.bond_yield].Value
                sp_ey:float = 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:float = 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:int = 6
        if len(self.market_data) >= min_count:
            market_closes:np.array = np.array([x[0] for x in self.market_data])
            market_returns:np.array = (market_closes[1:] - market_closes[:-1]) / market_closes[:-1]
            rf_rates:np.array = np.array([x[1] for x in self.market_data][1:])
            excess_returns:np.array = market_returns - rf_rates
            
            yield_gaps:List[float] = [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:float = yield_gaps[-1]
            
            # predicted market return
            Y:float = 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'