"""Bond returns and characteristics (TRACE only) from <https://openbondassetpricing.com/>"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/02_papers/dickerson_etal_2023.ipynb.

# %% ../../nbs/02_papers/dickerson_etal_2023.ipynb 4
from __future__ import annotations

import requests
import gzip
import zipfile
import io
import pandas as pd

from .. import wrds

# %% auto 0
__all__ = ['PROVIDER', 'URL', 'GZ_FILE', 'HOST_WEBSITE', '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', 'get_raw_data',
           'process_raw_data']

# %% ../../nbs/02_papers/dickerson_etal_2023.ipynb 5
PROVIDER = 'Alexander Dickerson, Philippe Mueller, Cesare Robotti, 2023'
URL = 'https://openbondassetpricing.com/wp-content/uploads/2023/10/WRDS_MMN_Corrected_Data.csv.zip' #contains a gzip file inside a zip file
GZ_FILE = 'WRDS_MMN_Corrected_Data.csv.gzip'
HOST_WEBSITE = 'https://openbondassetpricing.com/'
FREQ = 'M'
MIN_YEAR = 2002
MAX_YEAR = None
ENTITY_ID_IN_RAW_DSET = 'cusip' # 9 digit cusip
ENTITY_ID_IN_CLEAN_DSET = 'cusip'
TIME_VAR_IN_RAW_DSET = 'date'
TIME_VAR_IN_CLEAN_DSET = f'{FREQ}date'

# %% ../../nbs/02_papers/dickerson_etal_2023.ipynb 6
def get_raw_data(url: str=URL,
                 gz_file: str=GZ_FILE, # Name of the gzip file inside the zip file found at `url` 
            ) -> pd.DataFrame:
    """Download raw data from `url`"""

    response = requests.get(url)
    if response.status_code == 200:
        # Decompress the file first with zip and then with gzip
        with io.BytesIO(response.content) as compressed_file:
            with zipfile.ZipFile(compressed_file, 'r') as zip_ref:
                with zip_ref.open(gz_file) as gzipped_file:
                    with gzip.open(gzipped_file, 'rb') as final_gzip_file:
                        df = pd.read_csv(io.BytesIO(final_gzip_file.read()), header=0, index_col=0)
    else:
        print("Failed to download the file. Status code:", response.status_code)

    return df

# %% ../../nbs/02_papers/dickerson_etal_2023.ipynb 9
def process_raw_data(df: pd.DataFrame=None,
                     permno_to_bond_cusip: bool|pd.DataFrame=True, # Whether to download permno-cusip link. If DataFrame, must contain 'cusip'
                     ) -> pd.DataFrame:
    """Cleans up dates and optionally adds CRSP permnos"""

    df['dtdate'] = pd.to_datetime(df['date'])
    df['Mdate'] = df['dtdate'].dt.to_period('M')    

    if not permno_to_bond_cusip: return df
    if permno_to_bond_cusip is True: permno_to_bond_cusip = wrds.linking.bond_cusip_permno()

    df = df.merge(permno_to_bond_cusip, on=['cusip'], how='left')

    for col in ['permno','permco']:
        if col in df.columns: df[col] = df[col].astype('Int64').astype('category')

    return df 
