# Libraries
# --------------------------------------------------------------------------------------------------------
import pandas as pd
import numpy as np
import shap
import matplotlib.pyplot as plt
# --------------------------------------------------------------------------------------------------------

# Reading test data
# --------------------------------------------------------------------------------------------------------
def read_test_data(attribute_names):
    # Load test data
    X_test_pre = np.load('../gen_train_data/data/output/pre/X_test_pre.npy', allow_pickle=True)
    y_test_pre = np.load('../gen_train_data/data/output/pre/y_test_pre.npy', allow_pickle=True)
    X_test_post = np.load('../gen_train_data/data/output/post/X_test_post.npy', allow_pickle=True)
    y_test_post = np.load('../gen_train_data/data/output/post/y_test_post.npy', allow_pickle=True)

    # Type conversion needed    
    data_dic = {
        "X_test_pre": pd.DataFrame(X_test_pre, columns=attribute_names).convert_dtypes(),
        "y_test_pre": y_test_pre,
        "X_test_post": pd.DataFrame(X_test_post, columns=attribute_names).convert_dtypes(),
        "y_test_post": y_test_post,
    }
    return data_dic
# --------------------------------------------------------------------------------------------------------

if __name__ == "__main__":
    # Setup
    # --------------------------------------------------------------------------------------------------------
    # Retrieve attribute names in order
    attribute_names = list(np.load('../gen_train_data/data/output/attributes.npy', allow_pickle=True))
    # Reading data
    data_dic = read_test_data(attribute_names)
    method_names = {
        0: "ORIG",
        1: "ORIG_CW",
        2: "OVER",
        3: "UNDER"
    }
    model_choices = {
        "ORIG": "XGB",
        "ORIG_CW": "RF",
        "OVER": "XGB",
        "UNDER": "XGB"
    }
    # --------------------------------------------------------------------------------------------------------

    # Plot generation
    # --------------------------------------------------------------------------------------------------------
    for j, method in enumerate(['', '', 'over_', 'under_']):
        if method != 'over_':
            continue
        for i, group in enumerate(['pre', 'post']):
            X_test = data_dic['X_test_' + group]
            y_test = data_dic['y_test_' + group]
            print(f"{group}-{method_names[j]}")
            method_name = method_names[j]
            model_name = model_choices[method_name]
            shap_inter_vals = np.load(f'./output/shap_inter_values/{group}_{method_name}.npy')
            shap.summary_plot(shap_inter_vals, X_test, max_display=15)
            plt.savefig(f'./output/plots/inters/{group}_{method_name}_{model_name}.svg', format='svg', dpi=1250)