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

# %% auto 0
__all__ = ['PROVIDER', 'URL', 'get_series_info', 'default_raw_vars', 'parse_varlist', 'get_raw_data', 'process_raw_data',
           'search']

# %% ../../nbs/00_fred/00_fred.ipynb 4
from typing import List, Dict
import time

import pandas as pd

import pandasmore as pdm
from .fred_api import Fred

# %% ../../nbs/00_fred/00_fred.ipynb 5
PROVIDER = 'Federal Reserve Economic Data (FRED)'
URL = 'https://fred.stlouisfed.org/'

# %% ../../nbs/00_fred/00_fred.ipynb 6
def get_series_info(series: str, # 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"""
    
    return Fred(api_key=api_key).get_series_info(series)

# %% ../../nbs/00_fred/00_fred.ipynb 8
def default_raw_vars() -> pd.DataFrame:
    """List of FRED series that will be used in `get_raw_data()` if none are specified"""

    return ['TB3MS','DTB3','GS10','DGS10','GS1','DGS1','AAA','BAA','DAAA','DBAA','FEDFUNDS','DFF','CPIAUCSL','CPIAUCNS','INDPRO','IPB50001SQ','UNRATE','GDP','GDPC1','GNP','GNPC96','GDPPOT','USREC','RECPROUSM156N','CFNAI','UMCSENT','MICH','USEPUINDXM','USEPUNEWSINDXM','USEPUINDXD','VIXCLS','VXOCLS']


# %% ../../nbs/00_fred/00_fred.ipynb 10
def parse_varlist(vars: List[str]=None, #list of variables requested by user; if None, will use `default_raw_vars()`
                  api_key: str=None # FRED api key. If None, will use os.getenv("FRED_API_KEY")
                  ) -> Dict[str, list]:
    "Splits `vars` by frequency and returns dict with one pd.DataFrame per frequency"
     
    if vars is None: vars = default_raw_vars()

    # Collect frequency data for all series
    freq_store = []
    info_store = []
    for series in vars:
        info = get_series_info(series, api_key)
        info_store.append(info.to_frame().T)
        freq_store.append(pd.DataFrame({'series': [series], 'freq': [info['frequency_short']]}))

    info_store = pd.concat(info_store, ignore_index=True)
    freq_store = pd.concat(freq_store, ignore_index=True)

    # Split the series by frequency
    out = {'info': info_store}
    frequencies = list(freq_store['freq'].value_counts().index)
    for f in frequencies:
        out[f] = list(freq_store['series'].loc[freq_store.freq==f])

    return out

# %% ../../nbs/00_fred/00_fred.ipynb 16
def get_raw_data(vars: str=None, # FRED series name
               api_key: str=None # FRED api key. If None, will use os.getenv("FRED_API_KEY")
               ) -> Dict[str, pd.DataFrame]: 
    """Retrieves `vars` from FRED, splits them by frequency and returns dict with one pd.DataFrame per frequency"""

    api = Fred(api_key=api_key)

    varlist = parse_varlist(vars)

    out = {'info':varlist['info']}
    varlist.pop('info')
    for freq, vars in varlist.items():
        out[freq] = pd.concat([api.get_series(series).to_frame(name=series) for series in vars], axis=1).dropna(how='all')
        
    return out

# %% ../../nbs/00_fred/00_fred.ipynb 21
def process_raw_data(
        data: dict=None,  # keys are frequency str literals; same format as the output of `get_raw_data`
) -> pd.DataFrame: 
    """Processes data from FRED: cleans the date and sets it as index using `pdm.setup_tseries`"""

    out = {'info': data['info']}
    data.pop('info')  
    for freq, df in data.items():  
        df = df.reset_index().rename({'index':'date'},axis=1)
        df = pdm.setup_tseries(df, freq=freq).drop('date', axis=1)
        out[freq] = df 

    return out

# %% ../../nbs/00_fred/00_fred.ipynb 25
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`"""

    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].copy()
              
