# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os
import pickle
import gzip

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import spearmanr
from sklearn.preprocessing import MinMaxScaler
from tqdm import tqdm
# -

os.chdir('../..')

from src.data.constants import PUNCTUATION

pd.options.display.max_columns = None

# ## 1. LIME final results

# ### AG News

with open('data/external/latest_kirill/bert_lime_agnews_dev_500.pkl.gz', 'rb') as f:
    j = gzip.decompress(f.read())
results = pickle.loads(j)
print('loaded saved results', 'bert_lime_agnews_dev_500')

# +
vault = []

for i, info in tqdm(results.items()):
    for perturbation in info:
        pinfo = info[perturbation]
        if type(pinfo) != dict:
#            print(i, info, perturbation, pinfo)
            continue
        if 'spearman' not in pinfo:
            continue
        if 'orig_tokens' not in pinfo:
            continue
        vault += [dict(
            index=i,
            pinfo=pinfo,
            perturbation=perturbation,
            sentence_orig=pinfo['sentence_orig'],
            sentence=pinfo['sentence'],
            explanation=pinfo['exp'],
            topk=pinfo['top_5_intersect'],
            spearman=pinfo['spearman'],
        )]
# -

vault_df = pd.DataFrame(vault)
print(vault_df.shape)
vault_df.head()

# +
# vault_df.drop(columns=['sentence_orig', 'sentence', 'explanation'], inplace=True)
# -

kirill_results = pd.json_normalize(vault_df['pinfo'])
print(kirill_results.shape)
kirill_results.head()

kirill_results = pd.concat(
    (kirill_results, vault_df[['index', 'perturbation']]),
    axis=1
)
print(kirill_results.shape)
kirill_results.head()

kirill_results.drop(['exp', 'top_5_intersect', 'orig_indices', 'spearman', 'sentence', 'top_10_intersect'], axis=1, inplace=True)

# +
k = 5

kirill_results['top_k_idx_orig'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'orig_scores'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
kirill_results['top_k_idx_perturb'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores'])
        )[:k]
    ), 
    axis=1
)

kirill_results['k_intersection'] = kirill_results.apply(lambda x: len(x['top_k_idx_orig'].intersection(x['top_k_idx_perturb'])) / k, axis=1)

kirill_results['spearman'] = kirill_results.apply(lambda x: spearmanr(x['orig_scores'], x[f'perturb_scores'])[0], axis=1)
# -

comb = kirill_results.groupby('perturbation')['k_intersection'].mean().reset_index().merge(
    kirill_results.groupby('perturbation')['k_intersection'].std().reset_index(),
    on='perturbation'
)
comb['label']= np.round(comb['k_intersection_x'], 2).astype(str) + ' ± ' + np.round(comb['k_intersection_y'], 2).astype(str)
comb

comb = kirill_results.groupby('perturbation')['spearman'].mean().reset_index().merge(
    kirill_results.groupby('perturbation')['spearman'].std().reset_index(),
    on='perturbation'
)
comb['label']= np.round(comb['spearman_x'], 2).astype(str) + ' ± ' + np.round(comb['spearman_y'], 2).astype(str)
comb

kirill_results.groupby('perturbation')['k_intersection'].count()

# ### SST

with open('data/external/latest_kirill/bert_lime_sst_dev_500.pkl.gz', 'rb') as f:
    j = gzip.decompress(f.read())
results = pickle.loads(j)

# +
vault = []

for i, info in tqdm(results.items()):
    for perturbation in info:
        pinfo = info[perturbation]
        if type(pinfo) != dict:
#            print(i, info, perturbation, pinfo)
            continue
        if 'spearman' not in pinfo:
            continue
        if 'orig_tokens' not in pinfo:
            continue
        vault += [dict(
            index=i,
            pinfo=pinfo,
            perturbation=perturbation,
            sentence_orig=pinfo['sentence_orig'],
            sentence=pinfo['sentence'],
            explanation=pinfo['exp'],
            topk=pinfo['top_5_intersect'],
            spearman=pinfo['spearman'],
        )]
# -

vault_df = pd.DataFrame(vault)
print(vault_df.shape)
vault_df.head()

# +
# vault_df.drop(columns=['sentence_orig', 'sentence', 'explanation'], inplace=True)
# -

kirill_results = pd.json_normalize(vault_df['pinfo'])
print(kirill_results.shape)
kirill_results.head()

kirill_results = pd.concat(
    (kirill_results, vault_df[['index', 'perturbation']]),
    axis=1
)
print(kirill_results.shape)
kirill_results.head()

kirill_results.drop(['exp', 'top_5_intersect', 'orig_indices', 'spearman', 'sentence'], axis=1, inplace=True)

# +
k = 5

kirill_results['top_k_idx_orig'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'orig_scores'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
kirill_results['top_k_idx_perturb'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores'])
        )[:k]
    ), 
    axis=1
)

kirill_results['k_intersection'] = kirill_results.apply(lambda x: len(x['top_k_idx_orig'].intersection(x['top_k_idx_perturb'])) / k, axis=1)

kirill_results['spearman'] = kirill_results.apply(lambda x: spearmanr(x['orig_scores'], x[f'perturb_scores'])[0], axis=1)
# -

comb = kirill_results.groupby('perturbation')['k_intersection'].mean().reset_index().merge(
    kirill_results.groupby('perturbation')['k_intersection'].std().reset_index(),
    on='perturbation'
)
comb['label']= np.round(comb['k_intersection_x'], 2).astype(str) + ' ± ' + np.round(comb['k_intersection_y'], 2).astype(str)
comb

comb = kirill_results.groupby('perturbation')['spearman'].mean().reset_index().merge(
    kirill_results.groupby('perturbation')['spearman'].std().reset_index(),
    on='perturbation'
)
comb['label']= np.round(comb['spearman_x'], 2).astype(str) + ' ± ' + np.round(comb['spearman_y'], 2).astype(str)
comb

kirill_results.groupby('perturbation')['k_intersection'].count()

# ## 2. LIME-SHAP comparison

# ### AG News

# #### Kirill results

with open('data/external/latest_kirill/bert_lime_agnews_dev_500.pkl.gz', 'rb') as f:
    j = gzip.decompress(f.read())
results = pickle.loads(j)
print('loaded saved results', 'bert_lime_agnews_dev_500')

# +
vault = []

for i, info in tqdm(results.items()):
    for perturbation in info:
        pinfo = info[perturbation]
        if type(pinfo) != dict:
#            print(i, info, perturbation, pinfo)
            continue
        if 'spearman' not in pinfo:
            continue
        if 'orig_tokens' not in pinfo:
            continue
        vault += [dict(
            index=i,
            pinfo=pinfo,
            perturbation=perturbation,
            sentence_orig=pinfo['sentence_orig'],
            sentence=pinfo['sentence'],
            explanation=pinfo['exp'],
            topk=pinfo['top_5_intersect'],
            spearman=pinfo['spearman'],
        )]
# -

vault_df = pd.DataFrame(vault)
print(vault_df.shape)
vault_df.head()

# +
# vault_df.drop(columns=['sentence_orig', 'sentence', 'explanation'], inplace=True)
# -

kirill_results = pd.json_normalize(vault_df['pinfo'])
print(kirill_results.shape)
kirill_results.head()

kirill_results = pd.concat(
    (kirill_results, vault_df[['index', 'perturbation']]),
    axis=1
)
print(kirill_results.shape)
kirill_results.head()

kirill_results.drop(['exp', 'top_5_intersect', 'orig_indices', 'spearman', 'sentence', 'top_10_intersect'], axis=1, inplace=True)

# +
k = 5

kirill_results['top_k_idx_orig'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'orig_scores'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
kirill_results['top_k_idx_perturb'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores'])
        )[:k]
    ), 
    axis=1
)

kirill_results['k_intersection'] = kirill_results.apply(lambda x: len(x['top_k_idx_orig'].intersection(x['top_k_idx_perturb'])) / k, axis=1)

kirill_results['spearman'] = kirill_results.apply(lambda x: spearmanr(x['orig_scores'], x[f'perturb_scores'])[0], axis=1)
# -

# Get original indices

agnews_index_map = dict(zip(range(500), pd.read_parquet('data/external/agn_dev_500.parquet').index))

kirill_results['index'] = kirill_results['index'].map(agnews_index_map)

kirill_results.head()

# #### Steven results

steven_result_dict = pickle.load(open('data/processed/perturbation_df_dict_AGNews_BERT.pickle', 'rb'))

steven_result_dict.keys()

steven_results = pd.DataFrame()

for perturbation, df in steven_result_dict.items():

    df['index'] = df.index

    df['perturbation'] = perturbation

    df = df.copy()[['shap_orig_values', 'shap_orig_tokens', f'shap_{perturbation}_values', f'shap_{perturbation}_tokens', 'sentence', 'index', 'perturbation']]

    df.rename(columns={
        'shap_orig_values': 'orig_scores', 
        'shap_orig_tokens': 'orig_tokens', 
        f'shap_{perturbation}_values': 'perturb_scores', 
        f'shap_{perturbation}_tokens': 'perturb_tokens', 
        'sentence': 'sentence_orig', 
        'index': 'index', 
        'perturbation': 'perturbation'
    }, inplace=True)

    steven_results = steven_results.append(df)

steven_results.shape

# +
k = 5

steven_results['top_k_idx_orig'] = steven_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'orig_scores'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
steven_results['top_k_idx_perturb'] = steven_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores'])
        )[:k]
    ), 
    axis=1
)

steven_results['k_intersection'] = steven_results.apply(lambda x: len(x['top_k_idx_orig'].intersection(x['top_k_idx_perturb'])) / k, axis=1)

steven_results['spearman'] = steven_results.apply(lambda x: spearmanr(x['orig_scores'], x[f'perturb_scores'])[0], axis=1)
# -

# #### Merge results

comparison = kirill_results.merge(steven_results, on=['index', 'perturbation'], suffixes=['_LIME', '_SHAP'], how='inner')
comparison.head()


# +
# len(comparison.iloc[0]['perturb_tokens_LIME'])

# + tags=[]
# bert_tok_example = comparison.iloc[0]['perturb_tokens_SHAP']
# bert_tok_example

# + tags=[]
# bert_score_example = comparison.iloc[0]['perturb_scores_SHAP']
# bert_score_example

# +
# tok_score_df = pd.DataFrame({'tokens': bert_tok_example, 'scores': bert_score_example})
# tok_score_df.head()

# +
# tok_score_df['tokens'] = tok_score_df['tokens'].str.strip()

# +
# tok_score_df['is_punctuation'] = np.where(
#     tok_score_df['tokens'].isin(PUNCTUATION),
#     1,
#     0
# )

# +
# tok_score_df = tok_score_df[tok_score_df.copy()['is_punctuation'] == 0]

# +
# tok_score_df['is_subword'] = np.where(
#     tok_score_df['tokens'].str.contains('##'),
#     1,
#     0
# )

# +
# tok_score_df['is_subword'] = np.where(
#     tok_score_df['is_subword'].shift(-1) == 1,
#     1,
#     tok_score_df['is_subword']
# )

# +
# tok_score_df['id'] = tok_score_df.index

# +
# tok_score_df['id'] = np.where(
#     (tok_score_df['is_subword'].shift(1) == 1) & (tok_score_df['is_subword'] != 0),
#     None,
#     tok_score_df['id']
# )

# +
# tok_score_df['id'] = tok_score_df['id'].ffill()

# +
# tok_score_df['tokens'] = tok_score_df['tokens'].str.replace('##', '')

# +
# tok_score_df = tok_score_df.groupby('id').agg({'tokens': lambda x: ''.join(x), 'scores': 'sum'}).reset_index(drop=True)
# -

# - LIME doesn't output scores for punctuation
# - LIME preserves casing
# - LIME drops duplicates of a token
# - We adjust the SHAP tokens and scores to be consistent with LIME, allowing a comparison to be made
# - We then sort the tokens for each explainer so that a comparison can be made

def shap_adjustment(shap_token_list, shap_val_list):
    tok_score_df = pd.DataFrame(
        {'tokens': shap_token_list, 'scores': shap_val_list}
    )
    # Remove whitespace
    tok_score_df['tokens'] = tok_score_df['tokens'].str.strip()
    
    # Remove any punctuation tokens
    tok_score_df['is_punctuation'] = np.where(
        tok_score_df['tokens'].isin(PUNCTUATION),
        1,
        0
    )
    tok_score_df = tok_score_df[tok_score_df.copy()['is_punctuation'] == 0]
    
    # Merge subwords into orignal word and combine respective scores
    tok_score_df['is_subword'] = np.where(
        tok_score_df['tokens'].str.contains('##'),
        1,
        0
    )
    
    # Pick up start of word formed of subwords
    tok_score_df['is_subword'] = np.where(
        tok_score_df['is_subword'].shift(-1) == 1,
        1,
        tok_score_df['is_subword']
    )
    
    # Create unique ID for each word
    tok_score_df['id'] = tok_score_df.index
    
    tok_score_df['id'] = np.where(
        (tok_score_df['is_subword'].shift(1) == 1) & (tok_score_df['is_subword'] != 0),
        None,
        tok_score_df['id']
    )
    
    tok_score_df['id'] = tok_score_df['id'].ffill()
    
    # Remove ##s from subwords
    tok_score_df['tokens'] = tok_score_df['tokens'].str.replace('##', '')
    
    # Combine words
    tok_score_df = tok_score_df.groupby('id').agg(
        {
            'tokens': lambda x: ''.join(x), 
            'scores': 'sum'
        }
    ).reset_index(drop=True)
    
    # Take mean over repeated tokens (LIME drops duplicates)
    tok_score_df = tok_score_df.groupby('tokens').mean().reset_index()
    
    return list(tok_score_df['tokens']), list(tok_score_df['scores'])


# +
# Make token adjustments on perturbed SHAP outputs
adjusted_tokens = comparison.apply(
    lambda x: shap_adjustment(
        shap_token_list=x['perturb_tokens_SHAP'], 
        shap_val_list=x['perturb_scores_SHAP']
    ),
    axis=1
)

comparison['perturb_tokens_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[0])
comparison['perturb_scores_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[1])

# +
# Make token adjustments on original SHAP outputs
adjusted_tokens = comparison.apply(
    lambda x: shap_adjustment(
        shap_token_list=x['orig_tokens_SHAP'], 
        shap_val_list=x['orig_scores_SHAP']
    ),
    axis=1
)

comparison['orig_tokens_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[0])
comparison['orig_scores_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[1])
# -

# Filter examples where length of scores the ame between LIME and SHAP (may still be some inconsistencies in the implementations of the explainers)
new_comp = comparison.copy()[
    comparison['perturb_scores_LIME'].apply(len) == comparison['perturb_scores_SHAP_adj'].apply(len)
]

new_comp.head()

(new_comp['perturb_scores_SHAP_adj'].apply(len) == new_comp['orig_scores_SHAP_adj'].apply(len)).value_counts()

new_comp.shape

new_comp = new_comp[
    ['index', 'perturbation', 'orig_scores_LIME', 'orig_tokens_LIME', 'perturb_tokens_LIME', 'perturb_scores_LIME', 'orig_tokens_SHAP_adj', 'orig_scores_SHAP_adj', 'perturb_tokens_SHAP_adj', 'perturb_scores_SHAP_adj', 'k_intersection_SHAP', 'spearman_SHAP', 'k_intersection_LIME', 'spearman_LIME']
]

new_comp.head()

# Lower case LIME tokens to be consistent with SHAP
new_comp['perturb_tokens_LIME'] = new_comp.apply(lambda x: [w.lower() for w in x['perturb_tokens_LIME']], axis=1)
new_comp['orig_tokens_LIME'] = new_comp.apply(lambda x: [w.lower() for w in x['orig_tokens_LIME']], axis=1)

# Find idx order corresponding to ordered tokens for LIME and SHAP
new_comp['token_LIME_idx_order_perturb'] = new_comp.apply(lambda x: np.argsort(x['perturb_tokens_LIME']), axis=1)
new_comp['token_SHAP_idx_order_perturb'] = new_comp.apply(lambda x: np.argsort(x['perturb_tokens_SHAP_adj']), axis=1)

new_comp['perturb_tokens_LIME'] = new_comp.apply(lambda x: np.array(x['perturb_tokens_LIME'])[x['token_LIME_idx_order_perturb']], axis=1)
new_comp['perturb_scores_LIME'] = new_comp.apply(lambda x: np.array(x['perturb_scores_LIME'])[x['token_LIME_idx_order_perturb']], axis=1)

new_comp['perturb_tokens_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['perturb_tokens_SHAP_adj'])[x['token_SHAP_idx_order_perturb']], axis=1)
new_comp['perturb_scores_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['perturb_scores_SHAP_adj'])[x['token_SHAP_idx_order_perturb']], axis=1)

# Find idx order corresponding to ordered tokens for LIME and SHAP
new_comp['token_LIME_idx_order_orig'] = new_comp.apply(lambda x: np.argsort(x['orig_tokens_LIME']), axis=1)
new_comp['token_SHAP_idx_order_orig'] = new_comp.apply(lambda x: np.argsort(x['orig_tokens_SHAP_adj']), axis=1)

new_comp['orig_tokens_LIME'] = new_comp.apply(lambda x: np.array(x['orig_tokens_LIME'])[x['token_LIME_idx_order_orig']], axis=1)
new_comp['orig_scores_LIME'] = new_comp.apply(lambda x: np.array(x['orig_scores_LIME'])[x['token_LIME_idx_order_orig']], axis=1)

new_comp['orig_tokens_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['orig_tokens_SHAP_adj'])[x['token_SHAP_idx_order_orig']], axis=1)
new_comp['orig_scores_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['orig_scores_SHAP_adj'])[x['token_SHAP_idx_order_orig']], axis=1)

# ### Compute k-intersection and Spearman between explainers

# +
k = 5

new_comp['top_k_idx_LIME'] = new_comp.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores_LIME'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
new_comp['top_k_idx_SHAP'] = new_comp.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores_SHAP_adj'])
        )[:k]
    ), 
    axis=1
)

new_comp['k_intersection'] = new_comp.apply(lambda x: len(x['top_k_idx_LIME'].intersection(x['top_k_idx_SHAP'])) / k, axis=1)

new_comp['spearman'] = new_comp.apply(lambda x: spearmanr(x['perturb_scores_LIME'], x[f'perturb_scores_SHAP_adj'])[0], axis=1)
# -

new_comp.head()

# Filter non-name change perturbations - others are non-deterministic
new_comp = new_comp[new_comp['perturbation'].isin(['strip_punct', 'remove_commas', 'remove_all_punct', 'switch_gender'])]

new_comp.head()

new_comp.groupby('perturbation').count()['index']

new_comp['k_intersection'].mean()

new_comp['k_intersection'].std()

new_comp['spearman'].mean()

new_comp['spearman'].std()


def plot_examples(orig_tokens, orig_values, perturbation_tokens, perturbation_values, perturbation, return_df=False, savefig=False, filename=None):
    
    # SHAP df
    example_df_SHAP = pd.DataFrame({
        'shap_orig_tokens': orig_tokens[0],
        'shap_orig_values': orig_values[0],
        f'shap_{perturbation}_tokens': perturbation_tokens[0],
        f'shap_{perturbation}_values': perturbation_values[0]
    })

    # MinMax scale to get scores within 0,1
    example_df_SHAP['shap_orig_values'] = MinMaxScaler().fit_transform(np.abs(example_df_SHAP['shap_orig_values'].values.reshape(-1, 1))).flatten()
    example_df_SHAP[f'shap_{perturbation}_values'] = MinMaxScaler().fit_transform(np.abs(example_df_SHAP[f'shap_{perturbation}_values'].values.reshape(-1, 1))).flatten()

    # Sort by SHAP tokens before perturbation
    example_df_SHAP.sort_values('shap_orig_tokens', ascending=False, inplace=True)
    example_df_SHAP.reset_index(drop=True, inplace=True)

    # LIME df
    example_df_LIME = pd.DataFrame({
        'lime_orig_tokens': orig_tokens[1],
        'lime_orig_values': orig_values[1],
        f'lime_{perturbation}_tokens': perturbation_tokens[1],
        f'lime_{perturbation}_values': perturbation_values[1]
    })

    # MinMax scale to get scores within 0,1
    example_df_LIME['lime_orig_values'] = MinMaxScaler().fit_transform(np.abs(example_df_LIME['lime_orig_values'].values.reshape(-1, 1))).flatten()
    example_df_LIME[f'lime_{perturbation}_values'] = MinMaxScaler().fit_transform(np.abs(example_df_LIME[f'lime_{perturbation}_values'].values.reshape(-1, 1))).flatten()

    # Sort by SHAP tokens before perturbation
    example_df_LIME.sort_values('lime_orig_tokens', ascending=False, inplace=True)
    example_df_LIME.reset_index(drop=True, inplace=True)

    
    fig, axs = plt.subplots(2, 1, sharex=True, figsize=(12, 15))
    
    sns.barplot(x=example_df_SHAP.index, y='shap_orig_values', data=example_df_SHAP, ci=None, color='blue', alpha=0.3, label='Original', ax=axs[0])
    sns.barplot(x=example_df_SHAP.index, y=f'shap_{perturbation}_values', data=example_df_SHAP, ci=None, color='red', alpha=0.3, label=PERTURBATION_FULL_FORM_DICT[perturbation], ax=axs[0])

    sns.barplot(x=example_df_LIME.index, y='lime_orig_values', data=example_df_LIME, ci=None, color='blue', alpha=0.3, label='Original', ax=axs[1])
    sns.barplot(x=example_df_LIME.index, y=f'lime_{perturbation}_values', data=example_df_LIME, ci=None, color='red', alpha=0.3, label=PERTURBATION_FULL_FORM_DICT[perturbation], ax=axs[1])
    
    axs[0].set_xticklabels(example_df_SHAP['shap_orig_tokens'])
    axs[0].set_ylabel('Normalised SHAP score')
    axs[1].set_ylabel('Normalised LIME score')
    plt.tight_layout()
    plt.gcf().subplots_adjust(bottom=0.5)
    plt.xticks(rotation=90)
    axs[0].legend()
    axs[1].legend()
    
    if savefig:
        plt.savefig(f'reports/figures/bert_shap/{filename}', bbox_inches='tight')

    plt.show()

    if return_df:
        return example_df_SHAP, example_df_LIME

PERTURBATION_FULL_FORM_DICT = {
    'strip_punct': 'Strip end punctuation',
    'remove_commas': 'Remove commas',
    'remove_all_punct': 'Remove all punctuation',
    'change_first_name': 'Change first name',
    'change_last_name': 'Change last name',
    'change_location': 'Change location',
    'add_typo': 'Add typo',
    'switch_gender': 'Switch gender',
    'swap_adj': 'Swap adjective'
}

new_comp[(new_comp['k_intersection'] < 0.5) & (new_comp[['k_intersection_SHAP', 'k_intersection_LIME']].mean(axis=1) < 0.7)].sample(n=1)

# +
idx = 157

plot_examples(
    orig_tokens=(
        new_comp.loc[idx]['orig_tokens_SHAP_adj'], 
        new_comp.loc[idx]['orig_tokens_LIME']
    ), 
    orig_values=(
        new_comp.loc[idx]['orig_scores_SHAP_adj'], 
        new_comp.loc[idx]['orig_scores_LIME']
    ), 
    perturbation_tokens=(
        new_comp.loc[idx]['perturb_tokens_SHAP_adj'], 
        new_comp.loc[idx]['perturb_tokens_LIME']
    ), 
    perturbation_values=(
        new_comp.loc[idx]['perturb_scores_SHAP_adj'], 
        new_comp.loc[idx]['perturb_scores_LIME']
    ),
    perturbation='remove_all_punct',
    savefig=True,
    filename='SHAP_LIME_comp_AG.png'
)
# -

# ### SST

# #### Kirill results

with open('data/external/latest_kirill/bert_lime_sst_dev_500.pkl.gz', 'rb') as f:
    j = gzip.decompress(f.read())
results = pickle.loads(j)
print('loaded saved results', 'bert_lime_sst_dev_500')

# +
vault = []

for i, info in tqdm(results.items()):
    for perturbation in info:
        pinfo = info[perturbation]
        if type(pinfo) != dict:
#            print(i, info, perturbation, pinfo)
            continue
        if 'spearman' not in pinfo:
            continue
        if 'orig_tokens' not in pinfo:
            continue
        vault += [dict(
            index=i,
            pinfo=pinfo,
            perturbation=perturbation,
            sentence_orig=pinfo['sentence_orig'],
            sentence=pinfo['sentence'],
            explanation=pinfo['exp'],
            topk=pinfo['top_5_intersect'],
            spearman=pinfo['spearman'],
        )]
# -

vault_df = pd.DataFrame(vault)
print(vault_df.shape)
vault_df.head()

# +
# vault_df.drop(columns=['sentence_orig', 'sentence', 'explanation'], inplace=True)
# -

kirill_results = pd.json_normalize(vault_df['pinfo'])
print(kirill_results.shape)
kirill_results.head()

kirill_results = pd.concat(
    (kirill_results, vault_df[['index', 'perturbation']]),
    axis=1
)
print(kirill_results.shape)
kirill_results.head()

kirill_results.drop(['exp', 'top_5_intersect', 'orig_indices', 'spearman', 'sentence'], axis=1, inplace=True)

# +
k = 5

kirill_results['top_k_idx_orig'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'orig_scores'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
kirill_results['top_k_idx_perturb'] = kirill_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores'])
        )[:k]
    ), 
    axis=1
)

kirill_results['k_intersection'] = kirill_results.apply(lambda x: len(x['top_k_idx_orig'].intersection(x['top_k_idx_perturb'])) / k, axis=1)

kirill_results['spearman'] = kirill_results.apply(lambda x: spearmanr(x['orig_scores'], x[f'perturb_scores'])[0], axis=1)
# -

# Get original indices

agnews_index_map = dict(zip(range(500), pd.read_parquet('data/external/sst_dev_500.parquet').index))

kirill_results['index'] = kirill_results['index'].map(agnews_index_map)

kirill_results.head()

# #### Steven results

steven_result_dict = pickle.load(open('data/processed/perturbation_df_dict_SST_BERT.pickle', 'rb'))

steven_result_dict.keys()

steven_results = pd.DataFrame()

for perturbation, df in steven_result_dict.items():

    df['index'] = df.index

    df['perturbation'] = perturbation

    df = df.copy()[['shap_orig_values', 'shap_orig_tokens', f'shap_{perturbation}_values', f'shap_{perturbation}_tokens', 'sentence', 'index', 'perturbation']]

    df.rename(columns={
        'shap_orig_values': 'orig_scores', 
        'shap_orig_tokens': 'orig_tokens', 
        f'shap_{perturbation}_values': 'perturb_scores', 
        f'shap_{perturbation}_tokens': 'perturb_tokens', 
        'sentence': 'sentence_orig', 
        'index': 'index', 
        'perturbation': 'perturbation'
    }, inplace=True)

    steven_results = steven_results.append(df)

steven_results.shape

# +
k = 5

steven_results['top_k_idx_orig'] = steven_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'orig_scores'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
steven_results['top_k_idx_perturb'] = steven_results.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores'])
        )[:k]
    ), 
    axis=1
)

steven_results['k_intersection'] = steven_results.apply(lambda x: len(x['top_k_idx_orig'].intersection(x['top_k_idx_perturb'])) / k, axis=1)

steven_results['spearman'] = steven_results.apply(lambda x: spearmanr(x['orig_scores'], x[f'perturb_scores'])[0], axis=1)
# -

# #### Merge results

comparison = kirill_results.merge(steven_results, on=['index', 'perturbation'], suffixes=['_LIME', '_SHAP'], how='inner')
comparison.head()

comparison.shape

# - LIME doesn't output scores for punctuation
# - LIME preserves casing
# - LIME drops duplicates of a token
# - We adjust the SHAP tokens and scores to be consistent with LIME, allowing a comparison to be made
# - We then sort the tokens for each explainer so that a comparison can be made

# +
# Make token adjustments on perturbed SHAP outputs
adjusted_tokens = comparison.apply(
    lambda x: shap_adjustment(
        shap_token_list=x['perturb_tokens_SHAP'], 
        shap_val_list=x['perturb_scores_SHAP']
    ),
    axis=1
)

comparison['perturb_tokens_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[0])
comparison['perturb_scores_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[1])

# +
# Make token adjustments on original SHAP outputs
adjusted_tokens = comparison.apply(
    lambda x: shap_adjustment(
        shap_token_list=x['orig_tokens_SHAP'], 
        shap_val_list=x['orig_scores_SHAP']
    ),
    axis=1
)

comparison['orig_tokens_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[0])
comparison['orig_scores_SHAP_adj'] = adjusted_tokens.apply(lambda x: x[1])
# -

# Filter examples where length of scores the ame between LIME and SHAP (may still be some inconsistencies in the implementations of the explainers)
new_comp = comparison.copy()[
    comparison['perturb_scores_LIME'].apply(len) == comparison['perturb_scores_SHAP_adj'].apply(len)
]

new_comp.head()

(new_comp['perturb_scores_SHAP_adj'].apply(len) == new_comp['orig_scores_SHAP_adj'].apply(len)).value_counts()

new_comp = new_comp[new_comp['perturb_scores_SHAP_adj'].apply(len) == new_comp['orig_scores_SHAP_adj'].apply(len)]

new_comp.shape

new_comp = new_comp[
    ['index', 'perturbation', 'orig_scores_LIME', 'orig_tokens_LIME', 'perturb_tokens_LIME', 'perturb_scores_LIME', 'orig_tokens_SHAP_adj', 'orig_scores_SHAP_adj', 'perturb_tokens_SHAP_adj', 'perturb_scores_SHAP_adj', 'k_intersection_SHAP', 'spearman_SHAP', 'k_intersection_LIME', 'spearman_LIME']
]

new_comp.head()

# Lower case LIME tokens to be consistent with SHAP
new_comp['perturb_tokens_LIME'] = new_comp.apply(lambda x: [w.lower() for w in x['perturb_tokens_LIME']], axis=1)
new_comp['orig_tokens_LIME'] = new_comp.apply(lambda x: [w.lower() for w in x['orig_tokens_LIME']], axis=1)

# Find idx order corresponding to ordered tokens for LIME and SHAP
new_comp['token_LIME_idx_order_perturb'] = new_comp.apply(lambda x: np.argsort(x['perturb_tokens_LIME']), axis=1)
new_comp['token_SHAP_idx_order_perturb'] = new_comp.apply(lambda x: np.argsort(x['perturb_tokens_SHAP_adj']), axis=1)

new_comp['perturb_tokens_LIME'] = new_comp.apply(lambda x: np.array(x['perturb_tokens_LIME'])[x['token_LIME_idx_order_perturb']], axis=1)
new_comp['perturb_scores_LIME'] = new_comp.apply(lambda x: np.array(x['perturb_scores_LIME'])[x['token_LIME_idx_order_perturb']], axis=1)

new_comp['perturb_tokens_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['perturb_tokens_SHAP_adj'])[x['token_SHAP_idx_order_perturb']], axis=1)
new_comp['perturb_scores_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['perturb_scores_SHAP_adj'])[x['token_SHAP_idx_order_perturb']], axis=1)

# Find idx order corresponding to ordered tokens for LIME and SHAP
new_comp['token_LIME_idx_order_orig'] = new_comp.apply(lambda x: np.argsort(x['orig_tokens_LIME']), axis=1)
new_comp['token_SHAP_idx_order_orig'] = new_comp.apply(lambda x: np.argsort(x['orig_tokens_SHAP_adj']), axis=1)

new_comp['orig_tokens_LIME'] = new_comp.apply(lambda x: np.array(x['orig_tokens_LIME'])[x['token_LIME_idx_order_orig']], axis=1)
new_comp['orig_scores_LIME'] = new_comp.apply(lambda x: np.array(x['orig_scores_LIME'])[x['token_LIME_idx_order_orig']], axis=1)

new_comp['orig_tokens_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['orig_tokens_SHAP_adj'])[x['token_SHAP_idx_order_orig']], axis=1)
new_comp['orig_scores_SHAP_adj'] = new_comp.apply(lambda x: np.array(x['orig_scores_SHAP_adj'])[x['token_SHAP_idx_order_orig']], axis=1)

# ### Compute k-intersection and Spearman between explainers

# +
k = 5

new_comp['top_k_idx_LIME'] = new_comp.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores_LIME'])  # abs to ignore direction of importance value
        )[:k]
    ), 
    axis=1
)
new_comp['top_k_idx_SHAP'] = new_comp.apply(
    lambda x: 
    set(
        np.argsort(
            -np.abs(x[f'perturb_scores_SHAP_adj'])
        )[:k]
    ), 
    axis=1
)

new_comp['k_intersection'] = new_comp.apply(lambda x: len(x['top_k_idx_LIME'].intersection(x['top_k_idx_SHAP'])) / k, axis=1)

new_comp['spearman'] = new_comp.apply(lambda x: spearmanr(x['perturb_scores_LIME'], x[f'perturb_scores_SHAP_adj'])[0], axis=1)
# -

new_comp.head()

# Filter non-name change perturbations - others are non-deterministic
new_comp = new_comp[new_comp['perturbation'].isin(['strip_punct', 'remove_commas', 'remove_all_punct', 'switch_gender'])]

new_comp.head()

new_comp.groupby('perturbation').count()['index']

new_comp['k_intersection'].mean()

new_comp['k_intersection'].std()

new_comp['spearman'].mean()

new_comp['spearman'].std()

new_comp[(new_comp[['k_intersection_LIME', 'k_intersection_SHAP']].mean(axis=1) < 0.7) & (new_comp['k_intersection'] < 0.7)].sample(n=1)

# +
idx = 435

plot_examples(
    orig_tokens=(new_comp.loc[idx]['orig_tokens_SHAP_adj'], new_comp.loc[idx]['orig_tokens_LIME']), 
    orig_values=(new_comp.loc[idx]['orig_scores_SHAP_adj'], new_comp.loc[idx]['orig_scores_LIME']), 
    perturbation_tokens=(new_comp.loc[idx]['perturb_tokens_SHAP_adj'], new_comp.loc[idx]['perturb_tokens_LIME']), 
    perturbation_values=(new_comp.loc[idx]['perturb_scores_SHAP_adj'], new_comp.loc[idx]['perturb_scores_LIME']),
    perturbation='remove_commas',
    savefig=True,
    filename='SHAP_LIME_comp_SST.png'
)

StarcoderdataJupyterScriptsDedupFiltered