io.py

# -*- coding: utf-8 -*-
"""io.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1B79ZmyuEpCsmKNktWPlLOaD46ShrR-9c
"""

!pip install pandas matplotlib seaborn networkx scikit-learn plotly scipy statsmodels shap
# -*- coding: utf-8 -*-
"""Anxiety Intervention Analysis with Causal Mediation

This notebook adapts the MoE framework to incorporate causal mediation
analysis using the statsmodels library.  It aims to understand the *mechanism*
by which the intervention (group assignment) affects post-intervention anxiety,
specifically examining the mediating role of pre-intervention anxiety.

Workflow:
1. Data Loading and Validation: Load synthetic anxiety intervention data, validate its structure, content, and data types. Handle potential errors gracefully.
2. Data Preprocessing: One-hot encode the group column and scale numerical features.  Rename columns to be statsmodels-friendly.
3. Causal Mediation Analysis: Use statsmodels to perform causal mediation analysis, estimating direct and indirect effects.
4. SHAP Value Analysis: Quantify feature importance in the context of mediation.
5. Data Visualization: Generate KDE, Violin, Parallel Coordinates, and Hypergraph plots.
6. Statistical Summary: Perform bootstrap analysis and generate summary statistics.
7. LLM Insights Report: Synthesize findings using Grok, Claude, and Grok-Enhanced, emphasizing the mediation analysis.

Keywords: Causal Mediation, Mediation Analysis, Anxiety Intervention, statsmodels, LLMs, Explainability, SHAP, Data Visualization
"""

import warnings
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
import shap
import os
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import RandomForestRegressor
import numpy as np
from io import StringIO
import plotly.express as px
from scipy.stats import bootstrap
import statsmodels.formula.api as sm
import statsmodels.stats.mediation as sm_mediation  # Correct import

# Suppress specific warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, module="plotly")

# --- Constants ---
COLAB_ENV = False  # Assume not in Colab unless explicitly set
try:
    from google.colab import drive
    drive.mount("/content/drive")
    COLAB_ENV = True
except ImportError:
    print("Not running in Google Colab environment.")

OUTPUT_PATH = "./output_anxiety_causal_mediation/" if not COLAB_ENV else "/content/drive/MyDrive/output_anxiety_causal_mediation/"
PARTICIPANT_ID_COLUMN = "participant_id"
GROUP_COLUMN = "group"  # Original group column *before* one-hot encoding
ANXIETY_PRE_COLUMN = "anxiety_pre"
ANXIETY_POST_COLUMN = "anxiety_post"
MODEL_GROK_NAME = "grok-base"
MODEL_CLAUDE_NAME = "claude-3.7-sonnet"
MODEL_GROK_ENHANCED_NAME = "grok-enhanced"
LINE_WIDTH = 2.5
BOOTSTRAP_RESAMPLES = 500

# --- PLACEHOLDERS FOR API KEYS ---  (Replace with your actual API keys)
GROK_API_KEY = "YOUR_GROK_API_KEY"  # Replace with your actual Grok API key
CLAUDE_API_KEY = "YOUR_CLAUDE_API_KEY"  # Replace with your actual Claude API key

# --- Helper Functions ---

def create_output_directory(path):
    """Creates the output directory if it doesn't exist, handling errors."""
    try:
        os.makedirs(path, exist_ok=True)
        return True
    except OSError as e:
        print(f"Error creating output directory: {e}")
        return False

def load_data_from_synthetic_string(csv_string):
    """Loads data from a CSV string, handling errors."""
    try:
        csv_file = StringIO(csv_string)
        return pd.read_csv(csv_file)
    except pd.errors.ParserError as e:
        print(f"Error parsing CSV data: {e}")
        return None
    except Exception as e:
        print(f"Error loading data: {e}")
        return None

def validate_dataframe(df, required_columns):
    """Validates the DataFrame: checks for missing columns, non-numeric data,
    duplicate participant IDs, valid group labels, and plausible anxiety ranges.
    Returns True if valid, False otherwise.
    """
    if df is None:
        print("Error: DataFrame is None. Cannot validate.")
        return False

    missing_columns = [col for col in required_columns if col not in df.columns]
    if missing_columns:
        print(f"Error: Missing columns: {missing_columns}")
        return False

    for col in required_columns:
        if col != PARTICIPANT_ID_COLUMN and col != GROUP_COLUMN:
            if not pd.api.types.is_numeric_dtype(df[col]):
                print(f"Error: Non-numeric values found in column: {col}")
                return False

    if df[PARTICIPANT_ID_COLUMN].duplicated().any():
        print("Error: Duplicate participant IDs found.")
        return False

    valid_groups = ["Group A", "Group B", "Control"]  # Define valid group names
    invalid_groups = df[~df[GROUP_COLUMN].isin(valid_groups)][GROUP_COLUMN].unique()
    if invalid_groups.size > 0:
        print(f"Error: Invalid group labels found: {invalid_groups}")
        return False

    for col in [ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN]:
        if df[col].min() < 0 or df[col].max() > 10:
            print(f"Error: Anxiety scores in column '{col}' are out of range (0-10).")
            return False

    return True

def preprocess_data(df, group_column, anxiety_pre_column, anxiety_post_column):
    """Preprocesses the data: one-hot encodes the group column and scales numerical features.
       Also handles potential issues with column names for statsmodels.

    Args:
        df: The input DataFrame.
        group_column: Name of the group column.
        anxiety_pre_column: Name of the pre-intervention anxiety column.
        anxiety_post_column: Name of the post-intervention anxiety column.

    Returns:
        A tuple containing:
          - The preprocessed DataFrame.
          - A list of the one-hot encoded group column names.
          - The updated group_column name (in case of renaming).
          - The updated anxiety_pre_column name.
          - The updated anxiety_post_column name.
    """

    # 1.  Rename columns to be statsmodels-friendly (no spaces, special chars, etc.)
    df = df.rename(columns={
        PARTICIPANT_ID_COLUMN: "participant_id",
        group_column: "group_col",  # Use a consistent, safe name
        anxiety_pre_column: "anxiety_pre",
        anxiety_post_column: "anxiety_post"
    })
    group_column = "group_col"  # Update to the new name
    anxiety_pre_column = "anxiety_pre"
    anxiety_post_column = "anxiety_post"


    # 2. One-hot encode the group column (handle categorical variables)
    df = pd.get_dummies(df, columns=[group_column], prefix="group", drop_first=False) # Added drop_first=False
    encoded_group_cols = [col for col in df.columns if col.startswith("group_")]

    # 3. Scale numerical features
    numerical_cols = [anxiety_pre_column, anxiety_post_column] + encoded_group_cols
    scaler = MinMaxScaler()
    df[numerical_cols] = scaler.fit_transform(df[numerical_cols])

    return df, encoded_group_cols, group_column, anxiety_pre_column, anxiety_post_column

def analyze_text_with_llm(text, model_name):
    """Placeholder for LLM analysis.  Replace with actual API calls."""

    # --- Placeholder Implementation (Replace with Real API Calls) ---
    if model_name == MODEL_GROK_NAME:
        # Simulate Grok responses (replace with actual API call)
        if "causal mediation analysis" in text.lower():
            return "Grok-base: Causal mediation analysis explores intervention mechanisms, identifying how pre-anxiety mediates the relationship between group assignment and post-anxiety levels.  It suggests potential pathways of influence."
        elif "shap summary" in text.lower():
            return "Grok-base: SHAP values in the mediation context highlight the relative importance of pre-anxiety and group assignment in predicting post-anxiety, considering their mediating roles.  This helps understand feature contributions within the causal framework."
        else:
            return f"Grok-base: General analysis on '{text}'."

    elif model_name == MODEL_CLAUDE_NAME:
        # Simulate Claude responses (replace with actual API call)
        if "causal mediation analysis" in text.lower():
            return "Claude 3.7: Causal mediation analysis identifies mediating pathways, showing how pre-anxiety levels influence the intervention's effect on post-anxiety.  It quantifies the direct and indirect effects of the intervention."
        elif "parallel coordinates" in text.lower():
            return "Claude 3.7: Parallel coordinates visualize mediation effects across groups, showing individual trajectories and the influence of pre-anxiety as a mediator.  This allows for comparison of pathways across different groups."
        else:
            return f"Claude 3.7: Enhanced mediation analysis on '{text}'."

    elif model_name == MODEL_GROK_ENHANCED_NAME:
        # Simulate Grok-Enhanced responses (replace with actual API call)
        if "causal mediation analysis" in text.lower():
            return "Grok-Enhanced: Causal mediation analysis comprehensively reveals nuanced intervention mechanisms and mediating factors, specifically highlighting how pre-anxiety levels mediate the impact of the intervention on post-anxiety outcomes. It provides a detailed breakdown of direct and indirect effects."
        elif "hypergraph" in text.lower():
            return "Grok-Enhanced: Hypergraph visualizes participant clusters based on mediation pathways, showing how pre-anxiety levels and group assignments interact to influence post-anxiety, revealing potential subgroups with distinct mediation patterns."
        else:
            return f"Grok-Enhanced: In-depth mediation focused analysis on '{text}'."

    return f"Model '{model_name}' not supported."
    # --- End Placeholder ---

def perform_causal_mediation_analysis(df, treatment_column, mediator_column, outcome_column, output_path):
    """Performs causal mediation analysis using statsmodels, handling errors."""

    try:
        # Ensure treatment is categorical (using C() notation)
        med_formula = f"{mediator_column} ~ C({treatment_column})"
        med_model = sm.ols(formula=med_formula, data=df).fit()

        # Model for the outcome, including the treatment and the mediator
        out_formula = f"{outcome_column} ~ C({treatment_column}) + {mediator_column}"
        out_model = sm.ols(formula=out_formula, data=df).fit()

        # Causal mediation analysis (requires statsmodels >= 0.13)
        mediation_result = sm_mediation.Mediation(out_model, med_model, treatment_column, mediator_column).fit()

        # Format the output
        mediation_info = str(mediation_result.summary())

        print("Causal Mediation Analysis Results:\n", mediation_info)

        # Save the results to a file
        with open(os.path.join(output_path, "mediation_results.txt"), "w") as f:
            f.write(mediation_info)

        return mediation_info

    except Exception as e:
        error_message = f"Error during causal mediation analysis: {e}"
        print(error_message)
        return error_message

def calculate_shap_values(df, feature_columns, target_column, output_path):
    """Calculates and visualizes SHAP values, handling errors."""
    try:
        model_rf = RandomForestRegressor(random_state=42).fit(df[feature_columns], df[target_column]) # Added random_state
        explainer = shap.TreeExplainer(model_rf)
        shap_values = explainer.shap_values(df[feature_columns])

        plt.figure(figsize=(10, 8))
        plt.style.use('dark_background')
        shap.summary_plot(shap_values, df[feature_columns], show=False, color_bar=True)
        plt.savefig(os.path.join(output_path, 'shap_summary_mediation.png'))
        plt.close()
        return f"SHAP summary for features {feature_columns} predicting {target_column} (Mediation Context)"

    except Exception as e:
        print(f"Error during SHAP value calculation: {e}")
        return "Error: SHAP value calculation failed."

def create_kde_plot(df, column1, column2, output_path, colors):
    """Creates a KDE plot, handling errors."""
    try:
        plt.figure(figsize=(10, 6))
        plt.style.use('dark_background')
        sns.kdeplot(data=df[column1], color=colors[0], label=column1.capitalize(), linewidth=LINE_WIDTH)
        sns.kdeplot(data=df[column2], color=colors[1], label=column2.capitalize(), linewidth=LINE_WIDTH)
        plt.title('KDE Plot of Anxiety Levels (Mediation Analysis)', color='white')
        plt.legend(facecolor='black', edgecolor='white', labelcolor='white')
        plt.savefig(os.path.join(output_path, 'kde_plot_mediation.png'))
        plt.close()
        return f"KDE plot visualizing distributions of {column1} and {column2} (mediation analysis)"
    except KeyError as e:
        print(f"Error generating KDE plot: Column not found: {e}")
        return "Error: KDE plot generation failed.  Missing column."
    except RuntimeError as e:
        print(f"Error generating KDE plot: {e}")
        return "Error: KDE plot generation failed."
    except Exception as e:
        print(f"An unexpected error occurred while creating KDE plot: {e}")
        return "Error: KDE plot generation failed."

def create_violin_plot(df, group_column, y_column, output_path, colors):
    """Creates a violin plot, handling errors and one-hot encoded groups."""
    try:
        plt.figure(figsize=(10, 6))
        plt.style.use('dark_background')

        # Handling group column when already one-hot encoded
        encoded_group_cols = [col for col in df.columns if col.startswith(f"{group_column}_")]

        if len(encoded_group_cols) > 0:
            # Create a temporary column for group membership
            df['temp_group'] = np.nan
            for col in encoded_group_cols:
                group_name = col.split('_', 1)[1]  # Extract group name from encoded column
                df.loc[df[col] == 1, 'temp_group'] = group_name

            # Create violin plot
            sns.violinplot(data=df, x='temp_group', y=y_column, palette=colors[:len(encoded_group_cols)], linewidth=LINE_WIDTH)
            # Remove the temp group after plotting
            df.drop('temp_group', axis=1, inplace=True)
        else:
            # If group column is already categorical
            sns.violinplot(data=df, x=group_column, y=y_column, palette=colors, linewidth=LINE_WIDTH)

        plt.title('Violin Plot of Anxiety Distribution by Group (Mediation Analysis)', color='white')
        plt.savefig(os.path.join(output_path, 'violin_plot_mediation.png'))
        plt.close()
        return f"Violin plot showing {y_column} across groups (mediation analysis)"

    except KeyError as e:
        print(f"Error generating violin plot: Column not found: {e}")
        return "Error: Violin plot generation failed. Missing column."
    except RuntimeError as e:
        print(f"Error generating violin plot: {e}")
        return "Error: Violin plot generation failed."
    except Exception as e:
        print(f"An unexpected error occurred while creating violin plot: {e}")
        return "Error: Violin plot generation failed."

def create_parallel_coordinates_plot(df, group_column, anxiety_pre_column, anxiety_post_column, output_path, colors):
    """Creates a parallel coordinates plot, handling one-hot encoded groups and errors."""
    try:
        # Prepare data: Need original group names, not one-hot encoded.
        plot_df = df[[group_column, anxiety_pre_column, anxiety_post_column]].copy()

        # Create a color map for groups
        unique_groups = plot_df[group_column].unique()
        group_color_map = {group: colors[i % len(colors)] for i, group in enumerate(unique_groups)}

        # Map group names to colors
        plot_df['color'] = plot_df[group_column].map(group_color_map)

        # Create the parallel coordinates plot
        fig = px.parallel_coordinates(
            plot_df,
            color='color',  # Use the new 'color' column
            dimensions=[anxiety_pre_column, anxiety_post_column],
            title="Anxiety Levels: Pre- vs Post-Intervention by Group (Mediation Analysis)",
            color_continuous_scale=px.colors.sequential.Viridis, # Using Viridis
        )

        # Customize appearance
        fig.update_layout(
            plot_bgcolor='black',
            paper_bgcolor='black',
            font_color='white',
            title_font_size=16,
        )

        fig.write_image(os.path.join(output_path, 'parallel_coordinates_plot_mediation.png'))
        return f"Parallel coordinates plot of anxiety pre vs post intervention by group (mediation analysis)"

    except KeyError as e:
        print(f"Error generating parallel coordinates plot: Column not found: {e}")
        return "Error: Parallel coordinates plot generation failed. Missing column."
    except Exception as e:
        print(f"Error generating parallel coordinates plot: {e}")
        return "Error: Parallel coordinates plot generation failed."

def visualize_hypergraph(df, anxiety_pre_column, anxiety_post_column, output_path, colors):
    """Visualizes a hypergraph of participant relationships, handling errors."""
    try:
        G = nx.Graph()
        participant_ids = df[PARTICIPANT_ID_COLUMN].tolist()
        G.add_nodes_from(participant_ids, bipartite=0)

        # Use .loc for correct indexing and avoid SettingWithCopyWarning
        feature_sets = {
            "anxiety_pre": df.loc[df[anxiety_pre_column] > df[anxiety_pre_column].mean(), PARTICIPANT_ID_COLUMN].tolist(),
            "anxiety_post": df.loc[df[anxiety_post_column] > df[anxiety_post_column].mean(), PARTICIPANT_ID_COLUMN].tolist()
        }
        feature_nodes = list(feature_sets.keys())
        G.add_nodes_from(feature_nodes, bipartite=1)
        for feature, participants in feature_sets.items():
            for participant in participants:
                G.add_edge(participant, feature)
        pos = nx.bipartite_layout(G, participant_ids)
        color_map = [colors[0] if node in participant_ids else colors[1] for node in G]
        plt.figure(figsize=(12, 10))
        plt.style.use('dark_background')
        nx.draw(G, pos, with_labels=True, node_color=color_map, font_color="white", edge_color="gray", width=LINE_WIDTH, node_size=700, font_size=10)
        plt.title("Hypergraph Representation of Anxiety Patterns (Mediation Analysis)", color="white")
        plt.savefig(os.path.join(output_path, "hypergraph_mediation.png"))
        plt.close()
        return "Hypergraph visualizing participant relationships based on anxiety pre and post intervention (mediation analysis)"

    except KeyError as e:
        print(f"Error generating hypergraph: Column not found: {e}")
        return "Error: Hypergraph generation failed. Missing column."
    except Exception as e:
        print(f"Error creating hypergraph: {e}")
        return "Error creating hypergraph."

def perform_bootstrap(data, statistic, n_resamples=BOOTSTRAP_RESAMPLES):
    """Performs bootstrap analysis and returns confidence intervals, handling errors."""
    try:
        bootstrap_result = bootstrap((data,), statistic, n_resamples=n_resamples, method='percentile', random_state=42) # Added random_state
        return bootstrap_result.confidence_interval
    except Exception as e:
        print(f"Error during bootstrap analysis: {e}")
        return (None, None)

def save_summary(df, bootstrap_ci, causal_mediation_info, output_path):
    """Saves summary statistics and bootstrap CI to a text file, handling errors."""
    try:
        summary_text = df.describe().to_string() + f"\nBootstrap CI for anxiety_post mean: {bootstrap_ci}\n\nCausal Mediation Analysis Summary:\n{causal_mediation_info}"
        with open(os.path.join(output_path, 'summary.txt'), 'w') as f:
            f.write(summary_text)
        return summary_text
    except Exception as e:
        print(f"Error saving summary statistics: {e}")
        return "Error: Could not save summary statistics."

def generate_insights_report(summary_stats_text, shap_analysis_info, kde_plot_desc,
                             violin_plot_desc, parallel_coords_desc, hypergraph_desc,
                             causal_mediation_info, output_path):
    """Generates an insights report using LLMs (placeholders), handling errors."""
    try:
        grok_insights = (
            analyze_text_with_llm(f"Analyze summary statistics:\n{summary_stats_text}", MODEL_GROK_NAME) + "\n\n" +
            analyze_text_with_llm(f"Interpret Causal Mediation Analysis results:\n{causal_mediation_info}", MODEL_GROK_NAME) + "\n\n" +
            analyze_text_with_llm(f"Explain SHAP summary in mediation context: {shap_analysis_info}", MODEL_GROK_NAME) + "\n\n"
        )
        claude_insights = (
            analyze_text_with_llm(f"Interpret KDE plot in mediation analysis context: {kde_plot_desc}", MODEL_CLAUDE_NAME) + "\n\n" +
            analyze_text_with_llm(f"Interpret Violin plot in mediation analysis context: {violin_plot_desc}", MODEL_CLAUDE_NAME) + "\n\n" +
            analyze_text_with_llm(f"Interpret Parallel Coordinates Plot in mediation analysis context: {parallel_coords_desc}", MODEL_CLAUDE_NAME) + "\n\n" +
            analyze_text_with_llm(f"Interpret Hypergraph in mediation analysis context: {hypergraph_desc}", MODEL_CLAUDE_NAME) + "\n\n"
        )
        grok_enhanced_insights = analyze_text_with_llm(f"Provide enhanced insights on anxiety intervention mechanisms based on causal mediation analysis, SHAP, and visualizations, focusing on mediating role of pre-anxiety.", MODEL_GROK_ENHANCED_NAME)

        combined_insights = f"""
    Combined Insights Report: Anxiety Intervention Mechanism Analysis (Causal Mediation)

    Grok-base Analysis:
    {grok_insights}

    Claude 3.7 Sonnet Analysis:
    {claude_insights}

    Grok-Enhanced Analysis (Mediation Focused):
    {grok_enhanced_insights}

    Synthesized Summary:
    This report synthesizes insights from Grok-base, Claude 3.7 Sonnet, and Grok-Enhanced, focusing on causal mediation analysis to understand the mechanisms of anxiety intervention effectiveness. Grok-base provides a statistical overview and interprets the mediation analysis results, highlighting potential pathways of influence and the importance of pre-anxiety as a mediator. Claude 3.7 Sonnet details visual patterns and distributions, contextualized within the mediation analysis, showing group differences and the shift towards lower anxiety. Grok-Enhanced, with a mediation focus, delivers nuanced interpretations and actionable recommendations based on the causal mediation analysis, SHAP values, and visualizations. It emphasizes the mediating role of pre-anxiety in the intervention's impact and provides a detailed breakdown of direct and indirect effects. The combined expert analyses, enhanced by causal mediation techniques, provide a more mechanistic and comprehensive understanding of the intervention's effectiveness, revealing potential pathways of impact and informing targeted improvements to intervention design. The analysis suggests that the intervention's effect on post-anxiety is partly mediated by pre-anxiety levels.
    """
        with open(os.path.join(output_path, 'insights.txt'), 'w') as f:
            f.write(combined_insights)
        print(f"Insights saved to: {os.path.join(output_path, 'insights.txt')}")
        return "Insights report generated successfully."

    except Exception as e:
        print(f"Error generating insights report: {e}")
        return "Error generating insights report."

# --- Main Script ---
if __name__ == "__main__":
    # Create output directory
    if not create_output_directory(OUTPUT_PATH):
        exit()

    # Synthetic dataset (small, embedded in code)
    synthetic_dataset = """
participant_id,group,anxiety_pre,anxiety_post
P001,Group A,4,2
P002,Group A,3,1
P003,Group A,5,3
P004,Group B,6,5
P005,Group B,5,4
P006,Group B,7,6
P007,Control,3,3
P008,Control,4,4
P009,Control,2,2
P010,Control,5,5
"""
    # Load and validate data
    df = load_data_from_synthetic_string(synthetic_dataset)
    if df is None:
        exit()

    required_columns = [PARTICIPANT_ID_COLUMN, GROUP_COLUMN, ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN]
    if not validate_dataframe(df, required_columns):
        exit()

    # Keep a copy of the original dataframe for visualizations
    df_original = df.copy()

    # Preprocess data: One-hot encode and scale
    df, encoded_group_cols, group_column, anxiety_pre_column, anxiety_post_column = preprocess_data(df, GROUP_COLUMN, ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN)

    # Define columns for mediation analysis
    # Use the *safe* column names after preprocessing
    treatment_column = 'group_Group A'  # Example: Using 'Group A' as the treatment (one-hot encoded)
    mediator_column = 'anxiety_pre'
    outcome_column = 'anxiety_post'

    # Perform causal mediation analysis
    causal_mediation_info = perform_causal_mediation_analysis(df.copy(), treatment_column, mediator_column, outcome_column, OUTPUT_PATH)

    # SHAP analysis (using one-hot encoded columns)
    shap_feature_columns = encoded_group_cols + [anxiety_pre_column]
    shap_analysis_info = calculate_shap_values(df.copy(), shap_feature_columns, anxiety_post_column, OUTPUT_PATH)

    # Visualization colors
    neon_colors = ["#FF00FF", "#00FFFF", "#FFFF00", "#00FF00"]

    # Create visualizations (using the *original* DataFrame for plotting)
    kde_plot_desc = create_kde_plot(
        df_original, anxiety_pre_column, anxiety_post_column, OUTPUT_PATH, neon_colors[:2]
    )
    violin_plot_desc = create_violin_plot(
        df_original, GROUP_COLUMN, anxiety_post_column, OUTPUT_PATH, neon_colors
    )
    parallel_coords_desc = create_parallel_coordinates_plot(
        df_original, GROUP_COLUMN, anxiety_pre_column, anxiety_post_column, OUTPUT_PATH, neon_colors
    )
    hypergraph_desc = visualize_hypergraph(
        df_original, anxiety_pre_column, anxiety_post_column, OUTPUT_PATH, neon_colors[:2]
    )

    # Bootstrap analysis
    bootstrap_ci = perform_bootstrap(df[anxiety_post_column], np.mean)

    # Save summary statistics
    summary_stats_text = save_summary(df, bootstrap_ci, causal_mediation_info, OUTPUT_PATH)

    # Generate insights report
    generate_insights_report(summary_stats_text, shap_analysis_info, kde_plot_desc, violin_plot_desc, parallel_coords_desc, hypergraph_desc, causal_mediation_info, OUTPUT_PATH)

    print("Execution completed successfully - Causal Mediation Analysis Enhanced Notebook.")