# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_stability.ipynb.

# %% auto 0
__all__ = ['psi']

# %% ../nbs/01_stability.ipynb 3
import numpy as np
import pandas as pd
from numpy import ndarray

def psi(expected: ndarray, # expected array of values (reference)
        actual: ndarray, # actual array of values (production)
        buckets=10):
    "Calculate the PSI for a single numeric variable"

    df_expected = pd.DataFrame({'expected': expected})
    df_actual   = pd.DataFrame({'actual': actual}) 

    df_expected['expected_bins'], bins = pd.qcut(df_expected['expected'], q=buckets, retbins=True, labels=False, duplicates='drop')
    df_actual['actual_bins'] = pd.cut(df_actual['actual'], bins=bins, labels=False)

    psi_df = (
        df_expected['expected_bins'].value_counts(1).sort_index().reset_index().rename(columns={'index': 'bins', 'expected_bins': 'prop_expected'})
        .merge(df_actual['actual_bins'].value_counts(1).sort_index().reset_index().rename(columns={'index': 'bins', 'actual_bins': 'prop_actual'}),
               on='bins', how='left')
        .assign(psi = lambda df: np.where((df['prop_actual'] == 0) | (df['prop_expected'] == 0), 0, 
                                          (df['prop_actual'] - df['prop_expected'])*np.log(df['prop_actual']/df['prop_expected'])))
    )
    
    psi_value = psi_df['psi'].sum()
    return psi_value
