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

# %% ../../nbs/01_wrds/03_compa.ipynb 3
from __future__ import annotations
from pathlib import Path
from typing import List
import os

import pandas as pd
import numpy as np

import pandasmore as pdm
from . import wrds_api
from .. import RESOURCES

# %% auto 0
__all__ = ['raw_metadata', 'raw_metadata_extra', 'default_raw_vars', 'download', 'clean', 'book_equity', 'tobin_q',
           'issuance_vars', 'investment_vars', 'profitability_vars', 'cashflow_vars', 'liquidity_vars', 'leverage_vars',
           'payout_vars']

# %% ../../nbs/01_wrds/03_compa.ipynb 4
def raw_metadata(rawfile: str|Path=RESOURCES/'compa_variable_descriptions.csv', # location of the raw variable labels file
             ) -> pd.DataFrame:
    "Loads raw variable labels file, cleans it and returns it as a pd.DataFrame"

    df = pd.read_csv(rawfile)
    df['output_of'] = 'wrds.compa.clean'

    df['Variable Label'] = df.apply(lambda row: row['Description'].replace(row['Variable Name'].strip()+' -- ', ''), axis=1)
    df['Variable Label'] = df.apply(lambda row: row['Variable Label'].replace( '(' + row['Variable Name'].strip() + ')', ''), axis=1)
    df['Variable Name'] = df['Variable Name'].str.strip().str.lower()
    df = df[['Variable Name', 'Variable Label', 'output_of', 'Type']].copy()
    df.columns = ['name','label','output_of','type']
    return df

# %% ../../nbs/01_wrds/03_compa.ipynb 7
def raw_metadata_extra(wrds_username: str=None
             ) -> pd.DataFrame:
    "Collects metadata from WRDS `comp.funda` table and merges it with `variable_labels`."

    if wrds_username is None:
        wrds_username = os.getenv('WRDS_USERNAME')
        if wrds_username is None: wrds_username = input("Enter your WRDS username: ") 

    try:
        db = wrds_api.Connection(wrds_username = wrds_username)
        funda = db.describe_table('comp','funda')
        nr_rows = db.get_row_count('comp','funda')
    finally:
        db.close()

    meta = funda[['name','type']].copy()
    meta['nr_rows'] = nr_rows
    meta['wrds_library'] = 'comp'
    meta['wrds_table'] = 'funda'

    meta = meta.merge(raw_metadata()[['name','label']], how='left', on='name')
    
    meta['output_of'] = 'wrds.compa.download'
    meta = pdm.order_columns(meta,these_first=['name','label','output_of'])
    for v in list(meta.columns):
        meta[v] = meta[v].astype('string')
    
    return meta

# %% ../../nbs/01_wrds/03_compa.ipynb 9
def default_raw_vars():
    """Default variables used in `download` if none are specified."""

    return ['datadate', 'gvkey', 'cusip' ,'cik' ,'tic' ,'fyear' ,'fyr' ,'naicsh', 'sich' ,'exchg',  
            'lt' ,'at' ,'txditc' ,'pstkl' ,'pstkrv' ,'pstk' ,'csho' ,'ajex' , 'rdip',
            'act' ,'dvc' ,'xad','seq' ,'che' ,'lct' ,'dlc' ,'ib' ,'dvp' ,'txdi' ,'dp' ,
            'txp' ,'oancf' ,'ivncf' ,'fincf' ,'dltt' ,'mib','ceq' ,'invt' ,'cogs' , 'revt',
            'sale' ,'capx' ,'xrd' ,'txdb' ,'prcc_f' ,'sstk' ,'prstkc' ,'dltis' ,'dltr' ,'emp' ,
            'dd1' ,'ppegt' ,'ppent' ,'xint' ,'txt' ,'sppe' ,'gdwl' ,'xrent' ,'re' ,'dvpsx_f' ,
            'tstk' ,'wcap' ,'rect' ,'xsga' ,'aqc' ,'oibdp' ,'dpact' ,'fic' ,'ni' ,'ivao' ,'ivst' ,
            'dv' , 'intan' ,'pi' ,'txfo' ,'pifo' ,'xpp' ,'drc' ,'drlt' ,'ap' ,'xacc' ,'itcb']             

# %% ../../nbs/01_wrds/03_compa.ipynb 11
def download(vars: List[str]=None, # If None, downloads `default_raw_vars`; else `permno`, `permco`, and `date` are added by default
             obs_limit: int=None, #Number of rows to download. If None, full dataset will be downloaded
             wrds_username: str=None, #If None, looks for WRDS_USERNAME with `os.getenv`, then prompts you if needed
             start_date: str="01/01/1900", # Start date in MM/DD/YYYY format
             end_date: str=None #End date in MM/DD/YYYY format; if None, defaults to current date
             ) -> pd.DataFrame:
    """Downloads `vars` from `start_date` to `end_date` from WRDS `comp.funda` library and adds PERMNO and PERMCO as in CCM"""

    if vars is None: vars = default_raw_vars()
    vars = ','.join(['a.gvkey', 'a.datadate'] + 
                    [f'a.{x}' for x in vars if x not in ['datadate', 'gvkey']])

    limit_clause = f"LIMIT {obs_limit}" if obs_limit is not None else ""
    sql_string=f"""SELECT b.lpermno as permno, b.lpermco as permco, b.liid as iid, {vars}
                    FROM comp.funda AS a
                    INNER JOIN crsp.ccmxpf_lnkhist AS b ON a.gvkey = b.gvkey
                    WHERE datadate BETWEEN b.linkdt AND COALESCE(b.linkenddt, CURRENT_DATE)
                            AND b.linktype IN ('LU','LC') AND b.linkprim IN ('P','C')
                            AND indfmt='INDL' AND datafmt='STD' AND popsrc='D' AND consol='C'
                            AND datadate BETWEEN '{start_date}' AND COALESCE(%(end)s, CURRENT_DATE)
                    {limit_clause}
                """
    return wrds_api.download(sql_string, wrds_username=wrds_username, params={'end':end_date})

# %% ../../nbs/01_wrds/03_compa.ipynb 14
def clean(df: pd.DataFrame=None,        # If None, downloads `vars` using `download` function; else, must contain `permno` and `datadate` columns
          vars: List[str]=None,         # If None, downloads `default_raw_vars`
          obs_limit: int=None, #Number of rows to download. If None, full dataset will be downloaded
          wrds_username: str=None,      # If None, looks for WRDS_USERNAME with `os.getenv`, then prompts you if needed
          start_date: str="01/01/1900", # Start date in MM/DD/YYYY format
          end_date: str=None,           # End date. Default is current date          
          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`. If `df` is None, downloads `vars` using `download` function."""

    if df is None: df = download(vars=vars, obs_limit=obs_limit,  wrds_username=wrds_username, start_date=start_date, end_date=end_date)
    df = pdm.setup_panel(df, panel_ids='permno', time_var='datadate', freq='Y', **clean_kwargs)
    return df 

# %% ../../nbs/01_wrds/03_compa.ipynb 19
def book_equity(df: pd.DataFrame=None, # If None, downloads (and cleans) only required vars
                add_itcb=False,
                return_metadata: bool=False # If true, just returns metadata dictionary
                ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['at', 'lt', 'seq', 'ceq', 'txditc', 'pstk', 'pstkrv', 'pstkl', 'itcb']},
                'outputs': ['bookeq','shreq','pref_stock'],
                'labels': {'bookeq': 'Book equity', 'shreq': 'Shareholder equity', 'pref_stock': 'Preferred stock'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()

    df['pstk'] = df['pstk'].fillna(0)
    df['pref_stock'] = np.where(df['pstkrv'].isnull(), df['pstkl'], df['pstkrv'])
    df['pref_stock'] = np.where(df['pref_stock'].isnull(),df['pstk'], df['pref_stock'])

    df['shreq'] = np.where(df['seq'].isnull(), df['ceq'] + df['pstk'], df['seq'])
    df['shreq'] = np.where(df['shreq'].isnull(), df['at'] - df['lt'], df['shreq'])

    df['bookeq'] = df['shreq'] + df['txditc'].fillna(0) - df['pref_stock']
    if add_itcb: df['bookeq'] = df['bookeq'] + df['itcb'].fillna(0)
    
    return df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 26
def tobin_q(df: pd.DataFrame=None,      # If None, downloads (and cleans) only required vars
            return_metadata: bool=False # If True, just returns the metadata dictionary
            ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['at', 'lt', 'seq', 'ceq', 'txditc', 'pstk', 'pstkrv', 'pstkl', 'itcb','prcc_f','csho']},
                'outputs':  ['tobinq'],
                'labels': {'tobinq': 'Tobin Q'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()

    beq = book_equity(df)[['bookeq']].copy()
    df = df.join(beq)

    df['tobinq'] = (df['at'] - df['bookeq'] + df['prcc_f'] * df['csho']) / df['at']
    df = df.replace([np.inf, -np.inf], np.nan)
    return  df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 31
def issuance_vars(df: pd.DataFrame=None,        # If None, downloads (and cleans) only required vars
                    return_metadata: bool=False # If True, just returns the metadata dictionary
                    ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['at', 'lt', 'seq', 'ceq', 'txditc', 'pstk', 'pstkrv', 'pstkl', 'itcb',
                                                  'sstk','prstkc','dltis','dltr', 're', 'dlc','dltt']},
                'outputs': ['equityiss_tot','equityiss_cfs', 'debtiss_tot', 'debtiss_cfs', 'debtiss_bs'],
                'labels': {'equityiss_tot':'Equity issuance','equityiss_cfs':'Equity issuance', 
                           'debtiss_tot':'Debt issuance', 'debtiss_cfs':'Debt issuance', 'debtiss_bs':'Debt issuance'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()

    beq = book_equity(df)[['bookeq']].copy()
    df = df.join(beq)
    
    df['lag_at'] = pdm.lag(df['at'])

    df['equityiss_cfs'] = (df['sstk'].fillna(0) - df['prstkc'].fillna(0)) / df['lag_at']
    df['debtiss_cfs'] = (df['dltis'].fillna(0) - df['dltr'].fillna(0)) / df['lag_at']

    df['debtiss_bs'] = (pdm.rdiff(df['dltt']) + pdm.rdiff(df['dlc'].fillna(0))) / df['lag_at']

    df['equityiss_tot'] = (pdm.rdiff(df['bookeq']) - pdm.rdiff(df['re'])) / df['lag_at']
    df['debtiss_tot'] = (pdm.rdiff(df['at']) - pdm.rdiff(df['bookeq'])) / df['lag_at']
    
    df = df.replace([np.inf, -np.inf], np.nan)
    return df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 34
def investment_vars(df: pd.DataFrame=None,      # If None, downloads (and cleans) only required vars
                    return_metadata: bool=False # If True, just returns the metadata dictionary
                    ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['ppent','capx','at']},
                'outputs': ['ppentpch','capx2la'],
                'labels': {'ppentpch':'Pct change in net PPE','capx2la': 'CAPX to lagged assets'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()
    
    df['ppentpch'] = pdm.rpct_change(df['ppent'])
    df['capx2la'] = df['capx'] / pdm.lag(df['at'])
    df = df.replace([np.inf, -np.inf], np.nan)
    return df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 37
def profitability_vars(df: pd.DataFrame=None,      # If None, downloads (and cleans) only required vars
                        return_metadata: bool=False # If True, just returns the metadata dictionary
                        ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['ib','at']},
                'outputs': ['roa'],
                'labels': {'roa':'Return on assets'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()

    df['roa'] = df['ib'] / df['at']
    df = df.replace([np.inf, -np.inf], np.nan)
    return df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 40
def cashflow_vars(df: pd.DataFrame=None,      # If None, downloads (and cleans) only required vars
                        return_metadata: bool=False # If True, just returns the metadata dictionary
                        ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['dtdate','oancf','ib','dp','at']},
                'outputs': ['cflow2la_is', 'cflow2la_cfs', 'cflow2la_full'],
                'labels': {'cflow2la_is':'Cash flows to lagged assets', 
                           'cflow2la_cfs':'Cash flows to lagged assets', 
                           'cflow2la_full':'Cash flows to lagged assets'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()
    
    df['cflow2la_is'] = (df['ib']+df['dp']) / pdm.lag(df['at'])
    df['cflow2la_cfs'] = df['oancf'] / pdm.lag(df['at'])
    df['cflow2la_full'] = np.where(df.dtdate.dt.year<1987, df['cflow2la_is'], df['cflow2la_cfs'])

    df = df.replace([np.inf, -np.inf], np.nan)    
    return df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 43
def liquidity_vars(df: pd.DataFrame=None,      # If None, downloads (and cleans) only required vars
                        return_metadata: bool=False # If True, just returns the metadata dictionary
                        ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['che','at']},
                'outputs': ['cash2a'],
                'labels': {'cash2a':'Cash holdings to assets'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()

    df['cash2a'] = df['che'] / df['at']

    df = df.replace([np.inf, -np.inf], np.nan) 
    return df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 46
def leverage_vars(df: pd.DataFrame=None,      # If None, downloads (and cleans) only required vars
                        return_metadata: bool=False # If True, just returns the metadata dictionary
                        ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['dltt','dlc','at']},
                'outputs': ['booklev'],
                'labels': {'booklev':'Book leverage'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()

    df['booklev'] = (df['dltt'] + df['dlc']) / df['at']
    df.loc[df.booklev<0, 'booklev'] = 0
    df.loc[df.booklev>1, 'booklev'] = 1
    
    df = df.replace([np.inf, -np.inf], np.nan)         
    return df[metadata['outputs']].copy()

# %% ../../nbs/01_wrds/03_compa.ipynb 49
def payout_vars(df: pd.DataFrame=None,      # If None, downloads (and cleans) only required vars
                        return_metadata: bool=False # If True, just returns the metadata dictionary
                        ) -> pd.DataFrame:

    metadata = {'inputs': {'wrds.compa.clean': ['dvc','prstkc','at']},
                'outputs': ['div2la','rep2la'],
                'labels': {'div2la': 'Dividends to lagged assets',
                           'rep2la': 'Repurchases to lagged assets'}
    }      
    if return_metadata: return metadata

    reqs = metadata['inputs']['wrds.compa.clean']
    if df is None: df = clean(vars=reqs)
    df = df[reqs].copy()

    df['div2la'] = df['dvc'].fillna(0) / pdm.lag(df['at'])
    df['rep2la'] = df['prstkc'].fillna(0) / pdm.lag(df['at'])

    df = df.replace([np.inf, -np.inf], np.nan) 
    return df[metadata['outputs']].copy()
