# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_fred/00_fred.ipynb.

# %% auto 0
__all__ = ['download', 'clean', 'get_series_info', 'search', 'common_raw_vars']

# %% ../../nbs/00_fred/00_fred.ipynb 4
import os, time

import pandas as pd
from .fred_api import Fred

import pandasmore as pdm

# %% ../../nbs/00_fred/00_fred.ipynb 5
def download(series: str=None, # FRED series name
               api_key: str=None # FRED api key. If None, will use os.getenv("FRED_API_KEY")
               ) -> pd.DataFrame: 
    """Retrieves series from FRED and returns it as a pd.DataFrame."""

    if api_key is None: api_key = os.getenv("FRED_API_KEY")
    api = Fred(api_key=api_key)
    return api.get_series(series).to_frame(name=series)

# %% ../../nbs/00_fred/00_fred.ipynb 7
def clean(series: str=None, # FRED series name
               label: str=None, # Name you want to give to the series in the output DataFrame. If None, use lowercase of `series`
               api_key: str=None # FRED api key. If None, will use os.getenv("FRED_API_KEY")
               ) -> pd.DataFrame: 
    """Retrieves series from FRED, cleans the date and sets it as index"""

    if api_key is None: api_key = os.getenv("FRED_API_KEY")
    api = Fred(api_key=api_key)
    freq = api.get_series_info('GDP').frequency_short 
    time.sleep(1)
    if label is None: label = str.lower(series)
    df = download(series).rename(columns={series:label}).reset_index().rename({'index':'date'},axis=1)
    df = pdm.setup_tseries(df, freq=freq).drop('date', axis=1)
    return df 

# %% ../../nbs/00_fred/00_fred.ipynb 9
def get_series_info(series: str=None, # FRED series name
                    api_key: str=None # FRED api key. If None, will use os.getenv("FRED_API_KEY")
                    ) -> pd.Series:
    """Get metadata from FRED for given `series` from FRED"""

    if api_key is None: api_key = os.getenv("FRED_API_KEY")
    api = Fred(api_key=api_key)
    return api.get_series_info(series)

# %% ../../nbs/00_fred/00_fred.ipynb 11
def search(search_text: str=None, # What to search for
              order_by: str='popularity', # How to order search results; try `search_rank` if you don't find what you were looking for
              nr_results: int=10, # How many results to output
              api_key: str=None # FRED api key. If None, will use os.getenv("FRED_API_KEY")
              ) -> pd.DataFrame:
    """Search FRED for a given `search_text`, sort by popularity and return only the first `nr_results`"""

    if api_key is None: api_key = os.getenv("FRED_API_KEY")
    api = Fred(api_key=api_key)  
    return api.search(search_text, order_by=order_by)\
              .pipe(pdm.order_columns, ['title', 'popularity','frequency_short', 'observation_start', 'observation_end'])\
              .iloc[:nr_results]
              

# %% ../../nbs/00_fred/00_fred.ipynb 14
def common_raw_vars() -> pd.DataFrame:
    """List of FRED series commonly used in finance research"""

    vars = [
    ['TB3MS','yield_3mt','3-month treasury bill secondary market rate', 'M', 'NSA',1934],
    ['DTB3','yield_3mt','3-month treasury bill secondary market rate', 'D', 'NSA',1954],
    ['GS10','yield_10yt','10-year treasury constant maturity market yield', 'M', 'NSA',1953],
    ['DGS10','yield_10yt', '10-year treasury constant maturity market yield', 'D', 'NSA',1962],
    ['GS1','yield_1yt', '1-year treasury constant maturity market yield', 'M', 'NSA', 1953],
    ['DGS1','yield_1yt', '1-year treasury constant maturity market yield', 'D', 'NSA', 1962],
    ['AAA','yield_aaa', 'Moodys Aaa seasoned corporate bond yield', 'M', 'NSA', 1919], 
    ['BAA','yield_baa', 'Moodys Baa seasoned corporate bond yield', 'M', 'NSA', 1919], 
    ['DAAA','yield_aaa', 'Moodys Aaa seasoned corporate bond yield', 'D', 'NSA', 1983], 
    ['DBAA','yield_baa', 'Moodys Baa seasoned corporate bond yield', 'D', 'NSA', 1986],     
    ['FEDFUNDS','yield_fedf', 'federal funds effective rate', 'M', 'NSA', 1954], 
    ['DFF','yield_fedf', 'federal funds effective rate', 'D', 'NSA', 1954], 
    ['CPIAUCSL','cpi', 'cpi all urban consumers', 'M', 'SA', 1947],  
    ['CPIAUCNS','cpi', 'cpi all urban consumers', 'M', 'NSA', 1913],  
    ['INDPRO','indprod', 'industrial production', 'M', 'SA', 1919],
    ['IPB50001SQ','indprod', 'industrial production', 'Q', 'SA', 1919],
    ['UNRATE','unemp', 'unemployment rate', 'M', 'SA', 1948], 
    ['GDP','ngdp', 'nominal gdp','Q','SA',1947],
    ['GDPC1','rgdp', 'real gdp','Q','SA',1947],
    ['GNP','ngnp', 'nominal gnp','Q','SA',1947],
    ['GNPC96','rgnp', 'real gnp','Q','SA',1947],
    ['GDPPOT','pot_rgdp','real potential gdp','Q','SA',1949],
    ['USREC','rec_dum', 'nber recession indicator', 'M', 'NSA', 1854], 
    ['RECPROUSM156N','rec_prob', 'nber real time recession probability', 'M', 'NSA', 1967],
    ['CFNAI','cfnai','chicago fed national activity index','M','NSA',1967],
    ['UMCSENT','sent_mich','U Michigan consumer confidence index','M','NSA',1977],
    ['MICH','exp_inflation','U Michigan inflation expectation','M','NSA',1978],
    ['USEPUINDXM','pu_bbd','overall EPU index from Baker,Bloom,Davis','M', 'NSA',1985],
    ['USEPUNEWSINDXM','punews_bbd','news comp of EPU index from BBD','M', 'NSA',1985],
    ['USEPUINDXD','pu_bbd','D EPU from BBD','D','NSA',1985],
    ['VIXCLS','vix','VIX index on snp100 from CBOE 1month','D','NSA',1990],
    ['VXOCLS','vxo','VXO index on snp500 from CBOE','D','NSA',1986]
    ]
    return pd.DataFrame(data = vars, columns = ['id','label','desc','freq','season_adj','starts']) 

