"""Retrieve and process data from WRDS Bond Returns File"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/01_wrds/07_bondret.ipynb.

# %% ../../nbs/01_wrds/07_bondret.ipynb 3
from __future__ import annotations
from typing import List

import pandas as pd
import numpy as np

import pandasmore as pdm
from . import wrds_api

# %% auto 0
__all__ = ['PROVIDER', 'URL', 'LIBRARY', 'TABLE', 'LINK_TABLE', 'FREQ', 'MIN_YEAR', 'MAX_YEAR', 'ENTITY_ID_IN_RAW_DSET',
           'ENTITY_ID_IN_CLEAN_DSET', 'TIME_VAR_IN_RAW_DSET', 'TIME_VAR_IN_CLEAN_DSET', 'list_all_vars',
           'parse_varlist', 'get_raw_data', 'process_raw_data']

# %% ../../nbs/01_wrds/07_bondret.ipynb 4
PROVIDER = 'Wharton Research Data Services (WRDS)'
URL = 'https://wrds-www.wharton.upenn.edu/pages/get-data/wrds-bond-returns/'
LIBRARY = 'wrdsapps'
TABLE = 'bondret'
LINK_TABLE = 'bondcrsp_link'
FREQ = 'M'
MIN_YEAR = 2002
MAX_YEAR = None
ENTITY_ID_IN_RAW_DSET = 'permno'
ENTITY_ID_IN_CLEAN_DSET = 'permno'
TIME_VAR_IN_RAW_DSET = 'date'
TIME_VAR_IN_CLEAN_DSET = f'{FREQ}date'

# %% ../../nbs/01_wrds/07_bondret.ipynb 5
def list_all_vars() -> pd.DataFrame:
    "Collects names of all available variables from WRDS f`{LIBRARY}.{TABLE}`."

    try:
        db = wrds_api.Connection()
        df = db.describe_table(LIBRARY,TABLE).assign(wrds_library=LIBRARY, wrds_table=TABLE)
    finally:
        db.close()

    return df[['name','type','wrds_library','wrds_table']].copy()

# %% ../../nbs/01_wrds/07_bondret.ipynb 8
def parse_varlist(vars: List[str]=None,
                  required_vars = [],
                  ) -> str:
    """Validates that `vars` are available in `{LIBRARY}.{TABLE}` table and adds a. prefixes to variable names to feed into an SQL query"""

    # Get all available variables and add suffixes needed for the SQL query
    suffix_mapping = {TABLE: 'a.'}
    all_avail_vars = list_all_vars().drop_duplicates(subset='name',keep='first').copy()
    all_avail_vars['w_prefix'] = all_avail_vars.apply(lambda row: suffix_mapping[row['wrds_table']] + row['name'] , axis=1)

    if vars == '*' or vars is None: return ','.join(list(all_avail_vars['w_prefix']))
    
    # Add required vars to requested vars
    vars_to_get =  required_vars + [x for x in list(set(vars)) if x not in required_vars]

    # Validate variables to be downloaded (make sure that they are in the target database)
    invalid_vars = [v for v in vars_to_get if v not in list(all_avail_vars.name)]
    if invalid_vars: raise ValueError(f"These vars are not in the database: {invalid_vars}") 

    # Extract information on which variable comes from which wrds table, so we know what prefix to use
    vars_to_get = pd.DataFrame(vars_to_get, columns=['name'])
    get_these = vars_to_get.merge(all_avail_vars, how = 'left', on = 'name')
        
    return ','.join(list(get_these['w_prefix']))

# %% ../../nbs/01_wrds/07_bondret.ipynb 10
def get_raw_data(
        vars: List[str]='*', # Downloads all available variables by defaul
        required_vars: List[str]=['cusip','date'], #list of variables that will get downloaded, even if not in `vars`
        nrows: int=None, #Number of rows to download. If None, full dataset will be downloaded
        start_date: str=None, # Start date in MM/DD/YYYY format
        end_date: str=None #End date in MM/DD/YYYY format
) -> pd.DataFrame:
    """Downloads `vars` from `start_date` to `end_date` from WRDS `{LIBRARY}.{TABLE}` library and adds PERMNO and PERMCO"""
 
    wrds_api.validate_dates([start_date, end_date])
    vars = parse_varlist(vars, required_vars=required_vars)

    sql_string=f"""SELECT c.permno as permno, c.permco as permco, {vars}
                    FROM {LIBRARY}.{TABLE} AS a
                    INNER JOIN {LIBRARY}.{LINK_TABLE} AS c ON a.cusip = c.cusip 
                    WHERE date BETWEEN c.link_startdt AND COALESCE(c.link_enddt, CURRENT_DATE)
                """
    if start_date is not None: sql_string += r" AND date >= %(start_date)s"
    if end_date is not None: sql_string += r" AND date <= %(end_date)s"
    if nrows is not None: sql_string += r" LIMIT %(nrows)s"
    
    return wrds_api.download(sql_string,
                             params={'start_date':start_date, 'end_date':end_date, 'nrows':nrows})

# %% ../../nbs/01_wrds/07_bondret.ipynb 13
def process_raw_data(
        df: pd.DataFrame=None,  # Must contain `permno` and `datadate` columns   
        clean_kwargs: dict={},  # Params to pass to `pdm.setup_panel` other than `panel_ids`, `time_var`, and `freq`
) -> pd.DataFrame:
    """Applies `pandasmore.setup_panel` to `df`"""

    # Change some variables to categorical
    for col in ['cusip']:
        if col in df.columns:
            df[col] = df[col].astype('string').astype('category')

    # Set up panel structure
    df = pdm.setup_panel(df, 
                         drop_index_duplicates=False, # Multiple bonds (cusips) per permno per date
                         panel_ids=ENTITY_ID_IN_RAW_DSET, time_var=TIME_VAR_IN_RAW_DSET, freq=FREQ, 
                         panel_ids_toint=False, 
                         **clean_kwargs)
    return df 
