# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_core.ipynb (unless otherwise specified).

__all__ = ['is_numeric', 'drop_singletons', 'discretize', 'params', 'autolabel', 'clean_category']

# Cell

import pandas as pd
from pandas.api.types import is_numeric_dtype as isnum
from matplotlib.pyplot import rcParams

# Cell

def is_numeric(col:str):
    """Returns True iff already numeric, or can be coerced.
    Usage: df.apply(is_numeric)
    Usage: is_numeric(df['colname'])

    Returns Boolean series.

    From:
    https://stackoverflow.com/questions/54426845/how-to-check-if-a-pandas-dataframe-contains-only-numeric-column-wise

    """
    return isnum(col) or pd.to_numeric(col, errors='coerce').notnull().all()

# Cell

def drop_singletons(df, verbose=1) -> None:
    """Drop columns with < 2 unique values. Inplace."""
    dropcols = []
    for col in df:
        if len(df[col].unique()) < 2:
            df.drop(columns=col, inplace=True)
            dropcols.append(col)
    if verbose:
        print(f"  DROPPED {dropcols} because < 2 vals each.")

# Cell

def discretize(df, nbins=10, cut=pd.qcut,
               verbose=2, drop_useless=True):
    """Discretize columns in {df} to have at most {nbins} categories.
      * Categorical columns: take the Top n-1 plus "Other"
      * Continuous columns: cut into {nbins} using {cut}.

    Returns a new discretized dataframe with the same column names.
    Promotes discrete columns to categories.

    Parameters
    -----------
    df: Dataframe to discretize
    nbins: Max number of bins to use. May return fewer.
    cut: Cutting method. Default `pd.qcut`. Consider pd.cut, or write your own.
    verbose: 0: silent, 1: colnames, 2: (Default) top N for each column
    drop_useless: Removes columns that have < 2 unique values.

    Replaces numerical NA values with 'NA'.

    """
    out = pd.DataFrame(index = df.index)
    isnum = df.apply(is_numeric)
    for col in df:
        if verbose > 0:
            print(col, end=':\n\t')
        if isnum[col]:
            out[col] = cut(df[col], nbins, duplicates='drop')
        else:
            topN = df[col].value_counts(dropna=False).head(nbins).keys()
            # Promote to Category, add Other, Drop all but TopN
            out[col] = df[col].astype('category')
            try:
                out[col] = out[col].cat.add_categories(['Other'])
            except ValueError:
                pass  # Already had 'Other'
            out[col] = out[col].where(out[col].isin(topN), 'Other')
            out[col] = out[col].cat.remove_unused_categories()
        if verbose > 1:
            print('\n\t'.join(out[col].value_counts(dropna=False, sort=False)\
                              .to_string().split('\n')))
        elif verbose > 0:
            print()
    if drop_useless is True:
        drop_singletons(out)
    return out

# Cell

params = {'font.size': 12,
          'figure.titlesize': 'xx-large',
          'axes.titlesize': 'xx-large',
          'legend.title_fontsize': 'x-large',
          'legend.fontsize': 'medium',
          'axes.labelsize': 'large'}
rcParams.update(params)

# Cell

def autolabel(ax, border=False)->None:
    """Label bars in a barplot {ax} with their height.
    Thanks to matplotlib, composition.ai, and jsoma/chart.py.

    TODO: how to label with their legend labels?

    """
    ymin, ymax = ax.get_ylim()
    y_height = ymax - ymin
    for rect in ax.patches:
        height, width = rect.get_height(), rect.get_width()
        if not border:
            rect.set_ec(rect.get_fc())
        if height < 0.1 * y_height:
            continue
        x, y = rect.get_x(), rect.get_y()
        label_x, label_y = x + width/2, y + height/2
        label_text = int(height)
        ax.text(label_x, label_y, label_text, ha='center', va='center',
               clip_on=True, color='w', weight='bold', rotation=90)

# Cell

def clean_category(df, col:str)->None:
    """Remove unused categories from df.col, inplace.
    If not a category, do nothing.

    """
    try:
        old = len(df[col].cat.categories)
        df[col] = df[col].cat.remove_unused_categories()
        new = len(df[col].cat.categories)
        print(f'{col} compressed from {old} to {new} categories.')
    except AttributeError:
        # Not a category column
        pass