# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_plotting.ipynb.

# %% ../nbs/07_plotting.ipynb 3
from __future__ import annotations
from .import_essentials import *
import seaborn as sns
import matplotlib.pyplot as plt
from .evaluate import Explanation
from .data import TabularDataModule

# %% auto 0
__all__ = ['summary_plot', 'individual_plot']

# %% ../nbs/07_plotting.ipynb 4
def _process_exp(
    x: jnp.DeviceArray, # Dim: (N, k)
    cf: jnp.DeviceArray, # Dim: (N, k)
    pred_fn: Callable, # Predictive function
    return_valid: bool = False # Return only valid `x`, `cf`, `y_cf` or not
): # Return `x`, `cf`, `y_cf`
    """Process valid data."""
    
    assert x.shape == cf.shape
    assert len(x.shape) > 1

    y_cf = pred_fn(cf).round()
    
    if return_valid:
        y_x = pred_fn(x).round()
        idx = (y_x == y_cf).reshape(-1)
        x, cf, y_cf = x[idx], cf[idx], y_cf[idx]
    return x, cf, y_cf

# %% ../nbs/07_plotting.ipynb 6
def _diff(
    x: jnp.DeviceArray, # Dim: (N, k)
    cf: jnp.DeviceArray, # Dim: (N, k)
    pred_fn: Callable, # Predictive function
    datamodule: TabularDataModule, # Data module
    return_valid: bool = False, # Return only valid `x`, `cf`, `y_cf` or not
    return_outcome_col: bool = True # Return outcome column or not
) -> Tuple[pd.DataFrame, pd.DataFrame]: # Return continuous and categorical differences
    """Calculate differences between `x` and `cf`"""
    
    x, cf, y_cf = _process_exp(x, cf, pred_fn, return_valid)
    # wrap to numpy.array
    x, cf, y_cf = map(np.asarray, (x, cf, y_cf))
    x_inv, cf_inv = map(datamodule.inverse_transform, [x, cf], [y_cf, y_cf])
    cat_idx = datamodule.cat_idx
    cont_diff = pd.DataFrame(
        data=cf[:, :cat_idx] - x[:, :cat_idx],
        columns=datamodule._configs.continous_cols
    ) 
    cat_diff = (~cf_inv.iloc[:, cat_idx:-1].eq(x_inv.iloc[:, cat_idx:-1])).astype(float)
    if return_outcome_col:
        return pd.concat([cont_diff, cf_inv.iloc[:, -1:]], axis=1), \
            pd.concat([cat_diff, cf_inv.iloc[:, -1:]], axis=1)
    else:
        return cont_diff, cat_diff


# %% ../nbs/07_plotting.ipynb 7
def _swarmplot(
    cont_diff: pd.DataFrame, ax
):
    target_col = cont_diff.columns[-1]
    cont_diff = cont_diff.melt(
        id_vars=[target_col], value_name='Difference', var_name='Features'
    )
    sns.stripplot(
        y="Features",
        x="Difference",
        data=cont_diff,
        hue=target_col,
        alpha=0.5,
        size=5,
        jitter=0.3,
        ax=ax
    ).set(title="Summary Plot for Continuous Features")
    return ax

# %% ../nbs/07_plotting.ipynb 8
def _stripplot(
    cat_diff: pd.DataFrame, ax
):
    target_col = cat_diff.columns[-1]
    cat_diff = cat_diff.melt(
        id_vars=[target_col], value_name='Difference', var_name='Features'
    )
    # add noise for clear display 
    jitter = 0.3
    delta = np.random.uniform(-jitter/2, jitter/2, len(cat_diff['Difference']))
    cat_diff['Difference'] = cat_diff['Difference'].values + delta
    sns.stripplot(
        y="Features",
        x="Difference",
        data=cat_diff,
        hue=target_col,
        alpha=0.5,
        size=5,
        jitter=0.3,
        ax=ax
    ).set(title="Summary Plot for Categorical Features")
    return ax

# %% ../nbs/07_plotting.ipynb 9
def summary_plot(
    exp: Explanation, # Explanations to visualize
    sample_frac: float = 1.0, # Sample part of data for visualization
    only_valid: bool = False, # Use only valid data
    figsize=(15, 7) # Figure size
) -> plt.Figure:
    """Globally visualize generated explanations."""

    cont, cat = _diff(
        exp.X, exp.cfs, exp.pred_fn, exp.data_module, only_valid
    )
    fig, (ax1, ax2) = plt.subplots(1, 2)
    _swarmplot(cont.sample(frac=sample_frac), ax1)
    _stripplot(cat.sample(frac=sample_frac), ax2)

    # Resize to create overlap
    fig.set_size_inches(figsize[0], figsize[1])
    plt.subplots_adjust(wspace=.5)

    return fig

# %% ../nbs/07_plotting.ipynb 10
def _barplot(
    diff: pd.Series, ax
):
    sns.barplot(
        x=diff.values, y=diff.index, ax=ax
    ).set(title="Counterfactual Explanations for An Individual Sample")

# %% ../nbs/07_plotting.ipynb 11
def individual_plot(
    exp: Explanation, # Explanations to visualize
    idx: int, # Index for visualization
    figsize=(15, 7) # Figure plot
) -> plt.Figure:
    """Locally visualize individual explanations."""

    cont, cat = _diff(
        exp.X[idx, :].reshape(1, -1), exp.cfs[idx, :].reshape(1, -1), 
        exp.pred_fn, exp.data_module, return_outcome_col=False
    )
    diff = pd.concat([cont, cat], axis=1)

    # plot individual differences
    fig, ax = plt.subplots()
    _barplot(diff.iloc[0], ax)
    fig.set_size_inches(figsize[0], figsize[1])
    return fig
