import requests
from datetime import datetime as dt
import math
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()


def request_quandl(endpoint, **kwargs):
  params = {'order': 'asc', 'returns': 'numpy', 'authtoken': 'YOUR_AUTH_TOKEN'}
  params.update(kwargs)
  r = requests.get('https://www.quandl.com/api/v3/datasets/%s/data.json' % endpoint, params=params).json()
#  pp (r)
  return r['dataset_data']['data']

start_date = '1964-07-06'
collapse = 'weekly'

# get 10-Year Treasury Constant Maturity Rate:
r_dgs10 = request_quandl('FRED/DGS10', collapse=collapse, start_date=start_date)
# get 3-Month Treasury Bill: Secondary Market Rate:
r_dtb3 = request_quandl('FRED/DTB3', collapse=collapse, start_date=start_date)
# get NBER based Recession Indicators for the United States from the Peak through the Period preceding the Trough:
r_usrecdp = request_quandl('FRED/USRECDP', collapse=collapse, start_date=start_date)

df = {'datetime': [], 'spread': [], 'recession': [], 'probability': [], 'probability_ratio': []}
i = 0
while i < len(r_dtb3):
  date = dt.strptime(r_dtb3[i][0], '%Y-%m-%d')
  df['datetime'].append(date)
  spread = r_dgs10[i][1] - r_dtb3[i][1]
  df['spread'].append(float(spread))
  df['recession'].append(int(r_usrecdp[i][1] * 5))
  a = (100-(spread + 3)*25)
  if a > 0:
    probability = 3
    probability_ratio = a
  else:
    probability = 0
    probability_ratio = a
  df['probability'].append(probability)
  df['probability_ratio'].append(probability_ratio)
  i += 1

x = df['datetime']

plt.figure(1, (30, 20))

plt.subplot(211)
plt.title('Recession Probability')
plt.xlabel('Time')
plt.ylabel('Price')
plt.grid(True)
plt.plot(x, df['spread'])
plt.plot(x, df['recession'], color='green')
plt.plot(x, df['probability'], color='red')
plt.plot(x, [0] * len(x), color='black')
plt.legend(['Spread', 'Recession', 'Probability', 'Zero-Line'], loc='upper left')
plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
plt.gca().xaxis.set_major_locator(mdates.YearLocator())
plt.gcf().autofmt_xdate()

plt.subplot(212)
plt.title('Recession Probability')
plt.xlabel('Time')
plt.ylabel('Ratio')
plt.grid(True)
plt.plot(x, df['recession'], color='green')
plt.plot(x, df['probability_ratio'], color='red')
plt.plot(x, [0] * len(x), color='black')
plt.legend(['Recession', 'Probability', 'Zero-Line'], loc='upper left')
plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
plt.gca().xaxis.set_major_locator(mdates.YearLocator())
plt.gcf().autofmt_xdate()

plt.show()