"""
Tabular Feature Attribution Engine powered by SHAP.
Calculates and formats explicit attribution percentages for clinical features.
"""
import shap
import numpy as np
import pandas as pd

class TabularSHAPExplainer:
    def __init__(self, trained_model, baseline_background_data):
        self.model = trained_model
        # Use KernelExplainer for broad compatibility across different Keras structures
        self.explainer = shap.KernelExplainer(self.predict_wrapper, baseline_background_data)

    def predict_wrapper(self, data_array):
        # FIX: Dynamically handle verbose flag based on model type (Keras vs Scikit-Learn)
        try:
            return self.model.predict(data_array, verbose=0)
        except TypeError:
            # Fallback for Scikit-Learn models which don't accept a verbose parameter
            return self.model.predict(data_array)

    def compute_local_attributions(self, patient_profile_df):
        """
        Calculates exact risk contribution percentages for an individual patient profile.
        """
        shap_values = self.explainer.shap_values(patient_profile_df, silent=True)

        # Handle single-output matrix transformations cleanly
        if isinstance(shap_values, list):
            raw_shap = shap_values[0][0]
        else:
            raw_shap = shap_values[0] if len(shap_values.shape) > 1 else shap_values

        feature_contributions = []
        feature_names = patient_profile_df.columns.tolist()

        for index, name in enumerate(feature_names):
            attribution_score = float(raw_shap[index])
            # Scale raw scores to readable percentage points
            percentage_shift = attribution_score * 100
            feature_contributions.append({
                "Feature": name,
                "SHAP_Score": attribution_score,
                "Impact_Percentage": percentage_shift
            })

        return pd.DataFrame(feature_contributions)

if __name__ == "__main__":
    # Build functional validation blocks using a simple scikit-learn mock model
    from sklearn.linear_model import LogisticRegression

    mock_features = pd.DataFrame(np.random.randn(20, 3), columns=['HighBP', 'BMI', 'Smoker'])
    mock_labels = np.random.choice([0, 1], size=20)

    mock_clf = LogisticRegression().fit(mock_features, mock_labels)

    # Run standalone tests to verify the SHAP processing logic
    explainer = TabularSHAPExplainer(mock_clf, mock_features.mean().values.reshape(1, -1))
    single_patient = pd.DataFrame([[1.2, 0.4, -0.8]], columns=['HighBP', 'BMI', 'Smoker'])

    contributions_df = explainer.compute_local_attributions(single_patient)
    print("[SUCCESS] SHAP explainability layer loaded successfully:")
    print(contributions_df)