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

__all__ = ['no_missing_values', 'plot_hist', 'plot_feat_correlations', 'plot_feat_dendrogram']

# Cell
from forgebox.imports import *
import scipy
from scipy.cluster import hierarchy as hc
from typing import List, Union

# Cell
def no_missing_values(
    dataset: pd.DataFrame, missing_threshold: float = .6
) -> pd.DataFrame:
    """Find the features with a fraction of missing values above `missing_threshold`"""

    # Calculate the fraction of missing in each column
    missing_series = dataset.isnull().sum() / dataset.shape[0]
    missing_stats = pd.DataFrame(missing_series).rename(
        columns={'index': 'feature', 0: 'missing_fraction'})

    # Sort with highest number of missing values on top
    missing_stats = missing_stats.sort_values(
        'missing_fraction', ascending=False)
    display(missing_stats)

    # Find the columns with a missing percentage above the threshold
    record_missing = pd.DataFrame(
        missing_series[missing_series > missing_threshold])
    record_missing = record_missing.reset_index().rename(
        columns={'index': 'feature', 0: 'missing_fraction'})

    to_drop = list(record_missing['feature'])

    print(f'{len(to_drop)} features with greater than {missing_threshold*100:.1f}% missing values.')
    if to_drop:
        print('Removed features are:')
        for feat in to_drop:
            print(f"* {feat}")

    data = dataset.drop(columns=to_drop)
    return data

# Cell
def plot_hist(df: pd.DataFrame, feat2show: List[str]=None) -> None:
    """
    Plot histograms on columns feat2show
    default to plot all columns
    One historygram for each column
    """
    if feat2show is None: feat2show=df.columns
    nbr_plots = len(feat2show)
    ncols = 3
    if nbr_plots <= ncols:
        raise RuntimeError('Please feat2show requires at least four features')
    else:
        nrows=(nbr_plots-1) // ncols +1

    fig, axs = plt.subplots(nrows=nrows, ncols=ncols, figsize=(18, nrows * 4))
    for i, row in enumerate(axs):
        for j, ax in enumerate(row):
            idx = i*ncols+j
            if idx > nbr_plots-1: break
            f = feat2show[idx]
            ax.hist(df[f])
            ax.set_title(f)
    plt.show()

# Cell
def plot_feat_correlations(
    df: pd.DataFrame,
    as_triangle: bool=True,
    threshold: float=0.25,
    return_df: bool=False,
) -> Union[None, pd.DataFrame]:
    """Plot the pearson correlation between features as a heatmap

    Parameters:
        df:          The dataframe whose feature correlations will be computed
        as_triangle: When True (default value) only show the lower part for higher readibility
        threshold:   All correlations whose absolute value is smaller than threshold will not be plotted.
        return_df:   When True: return a dataframe with the full correlation matrix

    Return None or the correlation matrix depending on return_df
    """
    import scipy

    corr = df.corr()

    feats = corr.columns
    nfeats = len(feats)
    non_num_feats = set(df.columns).difference(set(corr.columns))

    if non_num_feats:
        print(f"Following features are not numerical and are excluded from correlation matrix: {', '.join(non_num_feats)}")

    if nfeats >= 25:
        print(f"There are {nfeats} features. The correlation matrix may not be easily readable in a notebook.")
        print("You can visualize the matrix as an image by right clicking on it, 'open image in a new tab, and zoom.")

    corr2plot = corr.values.copy()
    corr2plot[corr.abs() < threshold] = 0

    if as_triangle:
        corr2plot = np.tril(corr2plot)

    fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(nfeats//2, nfeats//2))
    cmap = 'seismic'
    im = ax.imshow(corr2plot, cmap=cmap, vmin=-1, vmax=+1)

    # Display each feature as a tick with label
    ax.set_xticks(np.arange(len(feats)))
    ax.set_xticklabels(feats, fontsize='small')
    ax.set_yticks(np.arange(len(feats)))
    ax.set_yticklabels(feats, fontsize='small')

    # Rotate x-ticks.
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")

    # Overlay each correlation values onto the heatmap
    for i in range(nfeats):
        for j in range(nfeats):
            text = ax.text(j, i, f"{corr2plot[i, j]:.2f}",
                           fontsize='x-small',
                           ha="center", va="center", color="w")

    new_line = '\n'
    title_l1 = 'Feature Correlation Matrix'
    title_l2 = f"Showing only correlations wider than +/-{threshold: .2f}"
    chart_title = f"{title_l1}{new_line}{title_l2}{f'{new_line}(Lower part only)' if as_triangle else ''}"
    ax.set_title(chart_title, fontsize='medium')

    fig.tight_layout()
    plt.show();

    if return_df:
        return corr

# Cell
def plot_feat_dendrogram(df: pd.DataFrame):
    """Plot the dendrograms of features based on spearman correlations

    Parameters:
        df: The dataframe whose feature correlations will be computed

    Return Nothing
    """
    corr = np.round(scipy.stats.spearmanr(df).correlation, 4)
    corr_condensed = hc.distance.squareform(1-corr)
    z = hc.linkage(corr_condensed, method='average')
    fig = plt.figure(figsize=(10,6))
    plt.title('Dendrogram of features based on spearman correlations', fontsize='medium')
    plt.xticks(fontsize='xx-small')
    hc.dendrogram(z, labels=df.columns, orientation='left', leaf_font_size='small')
    plt.show()