io.py

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

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1FupjvcVAvp_SUuobBwfs-6Twh0DpGPot
"""

# Refactored Code: Mix of Experts for Anxiety Intervention Explainability

"""
Robust Anxiety Intervention Analysis with Missing Data Handling and Mix of Experts Explainability

This script adapts the Mix of Experts (MoE) framework to analyze anxiety intervention data,
incorporating robust techniques for handling missing data and providing explainability
through visualizations, SHAP values, and insights from Large Language Models (LLMs).

It addresses missing data using iterative imputation, ensures data validation,
and generates comprehensive reports with visualizations and LLM-driven interpretations.

Key improvements include:
    - Enhanced error handling and informative messages.
    - Improved code readability and modularity.
    - Clearer function documentation and comments.
    - Robust handling of Plotly in Google Colab.
    - More descriptive variable names.
    - Consistent code style.
"""

import warnings
import os
from io import StringIO

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import networkx as nx
import shap
import plotly.express as px
from scipy.stats import bootstrap
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import RandomForestRegressor
from sklearn.experimental import enable_iterative_imputer  # noqa: F401
from sklearn.impute import IterativeImputer, KNNImputer
from sklearn.linear_model import BayesianRidge

# --- Configuration and Setup ---
# Suppress warnings (use with caution in production)
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning, module="plotly") # Corrected: category=UserWarning (class, not string)

# Plotly setup for Google Colab
PLOTLY_AVAILABLE = False
COLAB_ENV = False
try:
    import plotly.io as pio
    pio.renderers.default = "colab"
    PLOTLY_AVAILABLE = True
    from google.colab import drive
    drive.mount('/content/drive')
    COLAB_ENV = True
except ImportError:
    print("Plotly I/O module or Google Colab environment not detected. Visualization may be limited.")

# Constants - Project Configuration
OUTPUT_PATH = "./output_anxiety_analysis/" if not COLAB_ENV else "/content/drive/MyDrive/output_anxiety_analysis/"
PARTICIPANT_ID_COLUMN = "participant_id"
GROUP_COLUMN = "group"
ANXIETY_PRE_COLUMN = "anxiety_pre"
ANXIETY_POST_COLUMN = "anxiety_post"
MODEL_GROK_BASE_NAME = "grok-base"
MODEL_CLAUDE_SONNET_NAME = "claude-3.7-sonnet"
MODEL_GROK_ENHANCED_NAME = "grok-enhanced"
LINE_WIDTH = 2.5
BOOTSTRAP_RESAMPLES = 500
NEON_COLORS = ["#FF00FF", "#00FFFF", "#FFFF00", "#00FF00"] # Consistent color palette

# Placeholder API Keys - Replace with your actual API keys for LLMs
GROK_API_KEY = "YOUR_GROK_API_KEY"  # Placeholder - Security Warning
CLAUDE_API_KEY = "YOUR_CLAUDE_API_KEY"  # Placeholder - Security Warning

# --- Utility Functions ---
def create_directory(dir_path):
    """Creates a directory if it does not exist.

    Args:
        dir_path (str): The path to the directory to create.

    Returns:
        bool: True if directory creation was successful or directory already exists, False otherwise.
    """
    try:
        os.makedirs(dir_path, exist_ok=True)  # exist_ok=True prevents FileExistsError if dir exists
        return True
    except OSError as e:
        print(f"Error creating directory '{dir_path}': {e}")
        return False

def load_dataframe_from_csv_string(csv_string):
    """Loads a Pandas DataFrame from a CSV formatted string.

    Args:
        csv_string (str): CSV formatted data as a string.

    Returns:
        pd.DataFrame or None: DataFrame if loaded successfully, None otherwise.
    """
    try:
        csv_file_like_object = StringIO(csv_string)
        dataframe = pd.read_csv(csv_file_like_object)
        return dataframe
    except pd.errors.ParserError as e:
        print(f"CSV parsing error: {e}")
        return None
    except Exception as e:
        print(f"Error loading data: {e}")
        return None

def validate_dataframe_structure(dataframe, required_columns):
    """Validates the structure and content of the input DataFrame.

    Checks for missing columns, data types, duplicate IDs, valid group labels, and anxiety score ranges.

    Args:
        dataframe (pd.DataFrame): DataFrame to validate.
        required_columns (list): List of column names that must be present in the DataFrame.

    Returns:
        bool: True if DataFrame is valid, False otherwise.
    """
    if dataframe is None:
        print("Error: DataFrame is None and cannot be validated.")
        return False

    missing_cols = [col for col in required_columns if col not in dataframe.columns]
    if missing_cols:
        print(f"Error: DataFrame is missing required columns: {missing_cols}")
        return False

    for col in required_columns:
        if col not in [PARTICIPANT_ID_COLUMN, GROUP_COLUMN] and not pd.api.types.is_numeric_dtype(dataframe[col]):
            print(f"Error: Column '{col}' should be numeric but contains non-numeric data.")
            return False

    if dataframe[PARTICIPANT_ID_COLUMN].duplicated().any():
        print("Error: Duplicate participant IDs found in column '{PARTICIPANT_ID_COLUMN}'. IDs must be unique.")
        return False

    valid_groups = ["Group A", "Group B", "Control"]
    invalid_groups = dataframe[~dataframe[GROUP_COLUMN].isin(valid_groups)][GROUP_COLUMN].unique()
    if invalid_groups.size > 0:
        print(f"Error: Invalid group labels found: {invalid_groups}. Valid groups are: {valid_groups}")
        return False

    for anxiety_col in [ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN]:
        if not (0 <= dataframe[anxiety_col].min() and dataframe[anxiety_col].max() <= 10):
            print(f"Error: Anxiety scores in column '{anxiety_col}' are out of the valid range [0-10].")
            return False

    return True

def analyze_text_with_llm_placeholder(text, model_name):
    """Placeholder function for LLM-based text analysis.

    This function simulates calls to different LLMs (Grok-base, Claude 3.7 Sonnet, Grok-Enhanced)
    based on keywords in the input text. In a real application, this would be replaced with
    actual API calls to the respective LLM services.

    Args:
        text (str): Text to be analyzed.
        model_name (str): Name of the LLM model to simulate ("grok-base", "claude-3.7-sonnet", "grok-enhanced").

    Returns:
        str: Simulated LLM analysis result.
    """
    text_lower = text.lower()
    if model_name == MODEL_GROK_BASE_NAME:
        if "missing data handling" in text_lower:
            return "Grok-base: Iterative imputation is employed to handle missing data, ensuring a complete dataset for robust analysis."
        elif "shap summary" in text_lower:
            return "Grok-base: SHAP summary reveals feature importance on imputed data, highlighting influences on predicted anxiety levels."
        else:
            return f"Grok-base: General analysis of '{text}'."
    elif model_name == MODEL_CLAUDE_SONNET_NAME:
        if "missing data handling" in text_lower:
            return "Claude 3.7: Missing data techniques like iterative imputation minimize bias by estimating and filling in missing values based on data patterns."
        elif "violin plot" in text_lower:
            return "Claude 3.7: Violin plots on imputed data illustrate group distributions, enabling anxiety level comparisons across groups after addressing missing data."
        else:
            return f"Claude 3.7: Enhanced analysis with missing data handling for '{text}'."
    elif model_name == MODEL_GROK_ENHANCED_NAME:
        if "missing data handling" in text_lower:
            return "Grok-Enhanced: Iterative imputation provides a reliable and generalizable analysis, reducing bias from incomplete data and improving causal inference accuracy."
        elif "parallel coordinates" in text_lower:
            return "Grok-Enhanced: Parallel coordinates visualize anxiety trajectories on imputed data, robust to missingness, showing pre- to post-intervention changes per group."
        else:
            return f"Grok-Enhanced: In-depth robust analysis with missing data handling for '{text}'."
    return f"Model '{model_name}' is not supported in this placeholder."

def scale_dataframe_columns(dataframe, columns_to_scale):
    """Scales specified columns in a DataFrame using MinMaxScaler.

    Args:
        dataframe (pd.DataFrame): DataFrame containing columns to scale.
        columns_to_scale (list): List of column names to be scaled.

    Returns:
        pd.DataFrame or None: DataFrame with scaled columns if successful, None otherwise.
    """
    try:
        scaler = MinMaxScaler()
        dataframe[columns_to_scale] = scaler.fit_transform(dataframe[columns_to_scale])
        return dataframe
    except ValueError as e:
        print(f"Scaling error: Ensure columns '{columns_to_scale}' are numeric. Error details: {e}")
        return None
    except Exception as e:
        print(f"Unexpected error during scaling: {e}")
        return None

def handle_missing_data(dataframe, imputation_method="iterative", knn_neighbors=5):
    """Handles missing data in a DataFrame using various imputation methods.

    Supported methods: 'mean', 'median', 'most_frequent', 'knn', 'iterative'.

    Args:
        dataframe (pd.DataFrame): DataFrame with missing values to impute.
        imputation_method (str): Method for imputation. Defaults to 'iterative'.
        knn_neighbors (int): Number of neighbors for KNN imputation (if method='knn'). Defaults to 5.

    Returns:
        tuple: (imputed_dataframe, imputation_description) or (None, error_message).
               imputed_dataframe: DataFrame with imputed values. None if imputation fails.
               imputation_description: String describing the imputation method used, or error message.
    """
    if dataframe.empty:
        return None, "Error: Input DataFrame is empty."

    numeric_columns = dataframe.select_dtypes(include=np.number).columns
    if not numeric_columns.tolist(): # Check if list is empty
        return None, "Error: No numeric columns found in DataFrame for imputation."

    dataframe_imputed = dataframe.copy()
    imputation_description = ""

    try:
        if imputation_method == "mean":
            dataframe_imputed[numeric_columns] = dataframe_imputed[numeric_columns].fillna(dataframe[numeric_columns].mean())
            imputation_description = "Mean Imputation: Missing values replaced with the mean of each column."

        elif imputation_method == "median":
            dataframe_imputed[numeric_columns] = dataframe_imputed[numeric_columns].fillna(dataframe[numeric_columns].median())
            imputation_description = "Median Imputation: Missing values replaced with the median of each column."

        elif imputation_method == "most_frequent":
            for column in dataframe.columns: # Iterate over all columns to handle non-numeric as well if needed in future
                most_frequent_value = dataframe[column].mode()[0] if not dataframe[column].mode().empty else None # Handle case where mode is empty
                if most_frequent_value is not None: # Only fill if mode is found
                    dataframe_imputed[column] = dataframe_imputed[column].fillna(most_frequent_value)
            imputation_description = "Most Frequent Value Imputation: Missing values replaced with the most frequent value in each column."

        elif imputation_method == "knn":
            imputer_knn = KNNImputer(n_neighbors=knn_neighbors)
            dataframe_imputed[numeric_columns] = imputer_knn.fit_transform(dataframe[numeric_columns])
            imputation_description = f"KNN Imputation: Missing values imputed using KNN (k={knn_neighbors})."

        elif imputation_method == "iterative":
            imputer_iterative = IterativeImputer(estimator=BayesianRidge(), max_iter=10, random_state=42, n_nearest_features=None, imputation_order="ascending")
            dataframe_imputed[numeric_columns] = imputer_iterative.fit_transform(dataframe[numeric_columns])
            imputation_description = "Iterative Imputation: Missing values imputed using IterativeImputer with BayesianRidge estimator."

        else:
            raise ValueError(f"Invalid imputation method: '{imputation_method}'. Choose from: 'mean', 'median', 'most_frequent', 'knn', 'iterative'.")

        imputation_description += f"\nNumeric columns imputed: {', '.join(numeric_columns)}"
        initial_missing_rows = dataframe.isnull().any(axis=1).sum()
        final_missing_rows = dataframe_imputed.isnull().any(axis=1).sum()
        imputation_description += f"\nRows with missing values before imputation: {initial_missing_rows}"
        imputation_description += f"\nRows with missing values after imputation: {final_missing_rows}"

        if dataframe_imputed.isnull().any().any():
            print("Warning: Missing values still present after imputation. Review data and method.")
        else:
            print("Missing data imputation completed successfully.")

        return dataframe_imputed, imputation_description

    except ValueError as ve:
        print(f"ValueError during imputation: {ve}")
        return None, str(ve)
    except Exception as e:
        print(f"Unexpected error during imputation: {e}")
        return None, f"Imputation process failed: {e}"

def calculate_shap_and_visualize(dataframe, feature_columns, target_column, output_dir):
    """Calculates SHAP values and generates a SHAP summary plot.

    Args:
        dataframe (pd.DataFrame): DataFrame containing features and target.
        feature_columns (list): List of column names to be used as features.
        target_column (str): Name of the target column.
        output_dir (str): Directory to save the SHAP summary plot.

    Returns:
        str: Description of the SHAP analysis, or error message.
    """
    try:
        model_rf = RandomForestRegressor(random_state=42).fit(dataframe[feature_columns], dataframe[target_column])
        explainer_shap = shap.TreeExplainer(model_rf)
        shap_values = explainer_shap.shap_values(dataframe[feature_columns])

        plt.figure(figsize=(12, 8))
        plt.style.use("dark_background")
        shap.summary_plot(shap_values, dataframe[feature_columns], show=False, color_bar=True)
        plot_path = os.path.join(output_dir, "shap_summary_imputed.png")
        plt.savefig(plot_path)
        plt.close()
        return f"SHAP analysis completed. Summary plot saved to '{plot_path}'. Features: {feature_columns}, Target: {target_column}."

    except ValueError as ve:
        print(f"ValueError during SHAP calculation: {ve}")
        return f"Error: SHAP calculation failed due to ValueError: {ve}"
    except Exception as e:
        print(f"Error during SHAP calculation or visualization: {e}")
        return "Error: SHAP value calculation and visualization failed."

def create_kde_visualization(dataframe, column1, column2, output_dir, colors):
    """Generates and saves a Kernel Density Estimate (KDE) plot.

    Args:
        dataframe (pd.DataFrame): DataFrame containing the columns for KDE plot.
        column1 (str): Name of the first column for KDE plot.
        column2 (str): Name of the second column for KDE plot.
        output_dir (str): Directory to save the KDE plot.
        colors (list): List of colors for the KDE lines.

    Returns:
        str: Description of the KDE plot, or error message.
    """
    try:
        plt.figure(figsize=(10, 6))
        plt.style.use('dark_background')
        sns.kdeplot(data=dataframe[column1], color=colors[0], label=column1.capitalize(), linewidth=LINE_WIDTH)
        sns.kdeplot(data=dataframe[column2], color=colors[1], label=column2.capitalize(), linewidth=LINE_WIDTH)
        plt.title("KDE Plot of Anxiety Levels (Imputed Data)", color="white")
        plt.xlabel("Anxiety Level (Scaled)", color="white") # Added x-label
        plt.ylabel("Density", color="white") # Added y-label
        plt.legend(facecolor="black", edgecolor="white", labelcolor="white")
        plot_path = os.path.join(output_dir, "kde_plot_imputed.png")
        plt.savefig(plot_path)
        plt.close()
        return f"KDE plot generated and saved to '{plot_path}'. Columns: {column1}, {column2}."

    except KeyError as ke:
        print(f"KeyError in KDE plot creation: Column '{ke}' not found in DataFrame.")
        return f"Error: KDE plot generation failed. Column not found: {ke}."
    except Exception as e:
        print(f"Error during KDE plot generation: {e}")
        return "Error: KDE plot generation failed."

def create_violin_visualizations_by_group(dataframe, group_column_prefix, y_column, output_dir, colors):
    """Generates and saves violin plots for each group in one-hot encoded group columns.

    Args:
        dataframe (pd.DataFrame): DataFrame containing group columns and y_column.
        group_column_prefix (str): Prefix of the one-hot encoded group columns (e.g., 'group_').
        y_column (str): Column name for the y-axis of the violin plots.
        output_dir (str): Directory to save the violin plots.
        colors (list): List of colors for the violin plots.

    Returns:
        str: Descriptions of the generated violin plots, or error message.
    """
    plot_descriptions = []
    try:
        for group_col in dataframe.columns:
            if group_col.startswith(group_column_prefix):
                plt.figure(figsize=(10, 6))
                plt.style.use('dark_background')
                group_name = group_col.replace(group_column_prefix, "")
                group_dataframe = dataframe[dataframe[group_col] == 1]
                sns.violinplot(data=group_dataframe, y=y_column, color=colors[0], linewidth=LINE_WIDTH) # Using single color from palette
                plt.title(f"Violin Plot of {y_column.capitalize()} Distribution for {group_name} (Imputed Data)", color="white")
                plt.ylabel(y_column.capitalize(), color="white") # Added y-label
                plot_path = os.path.join(output_dir, f"violin_plot_{group_name}_imputed.png")
                plt.savefig(plot_path)
                plt.close()
                plot_descriptions.append(f"Violin plot for group '{group_name}' saved to '{plot_path}'. Y-column: {y_column}.")
        return "\n".join(plot_descriptions)

    except KeyError as ke:
        print(f"KeyError in violin plot creation: Column '{ke}' not found in DataFrame.")
        return f"Error: Violin plot generation failed. Column not found: {ke}."
    except Exception as e:
        print(f"Error during violin plot generation: {e}")
        return "Error: Violin plot generation failed."

def create_parallel_coordinates_visualization(dataframe, group_column_prefix, anxiety_pre_col, anxiety_post_col, output_dir, colors):
    """Generates and displays parallel coordinates plots for each group using Plotly.

    Args:
        dataframe (pd.DataFrame): DataFrame containing group columns and anxiety columns.
        group_column_prefix (str): Prefix of the one-hot encoded group columns.
        anxiety_pre_col (str): Column name for pre-intervention anxiety levels.
        anxiety_post_col (str): Column name for post-intervention anxiety levels.
        output_dir (str): Directory to save plots (not used for display in Colab).
        colors (list): List of colors (not directly used in Plotly, but kept for consistency).

    Returns:
        str: Descriptions of the generated parallel coordinates plots, or error message.
    """
    plot_descriptions = []
    for group_col in dataframe.columns:
        if group_col.startswith(group_column_prefix):
            group_name = group_col.replace(group_column_prefix, "")
            group_dataframe = dataframe[dataframe[group_col] == 1]
            plot_dataframe = group_dataframe[[anxiety_pre_col, anxiety_post_col]].copy()

            fig = px.parallel_coordinates(
                plot_dataframe,
                dimensions=[anxiety_pre_col, anxiety_post_col],
                title=f"Anxiety Levels: Pre- vs Post-Intervention for {group_name} (Imputed Data)",
                color_continuous_scale=px.colors.sequential.Viridis,
            )
            fig.update_layout(plot_bgcolor="black", paper_bgcolor="black", font_color="white", title_font_size=16)

            if COLAB_ENV and PLOTLY_AVAILABLE:
                print(f"Displaying parallel coordinates plot for {group_name} in Colab.")
                fig.show() # Display in Colab
            else:
                plot_path = os.path.join(output_dir, f"parallel_coords_{group_name}_imputed.png")
                fig.write_image(plot_path) # Save if not in Colab or Plotly not available
                plot_descriptions.append(f"Parallel coordinates plot for group '{group_name}' saved to '{plot_path}'.")
    return "\n".join(plot_descriptions)

def create_hypergraph_visualization(dataframe, anxiety_pre_col, anxiety_post_col, output_dir, colors):
    """Generates and saves a hypergraph visualization of anxiety patterns.

    Args:
        dataframe (pd.DataFrame): DataFrame containing participant IDs and anxiety levels.
        anxiety_pre_col (str): Column name for pre-intervention anxiety levels.
        anxiety_post_col (str): Column name for post-intervention anxiety levels.
        output_dir (str): Directory to save the hypergraph plot.
        colors (list): List of colors for nodes in the hypergraph.

    Returns:
        str: Description of the hypergraph visualization, or error message.
    """
    try:
        hypergraph = nx.Graph()
        participant_ids = dataframe[PARTICIPANT_ID_COLUMN].tolist()
        hypergraph.add_nodes_from(participant_ids, bipartite=0) # Participants as one set of nodes
        feature_sets = {
            "anxiety_pre_high": dataframe[PARTICIPANT_ID_COLUMN][dataframe[anxiety_pre_col] > dataframe[anxiety_pre_col].mean()].tolist(),
            "anxiety_post_high": dataframe[PARTICIPANT_ID_COLUMN][dataframe[anxiety_post_col] > dataframe[anxiety_post_col].mean()].tolist(),
        }
        feature_nodes = list(feature_sets.keys())
        hypergraph.add_nodes_from(feature_nodes, bipartite=1) # Features as the other set of nodes
        for feature, participants in feature_sets.items():
            for participant in participants:
                hypergraph.add_edge(participant, feature)

        node_positions = nx.bipartite_layout(hypergraph, participant_ids)
        node_colors = [colors[0] if node in participant_ids else colors[1] for node in hypergraph]

        plt.figure(figsize=(12, 10))
        plt.style.use('dark_background')
        nx.draw(hypergraph, node_positions, with_labels=True, node_color=node_colors, font_color="white", edge_color="gray", width=LINE_WIDTH, node_size=700, font_size=10)
        plt.title("Hypergraph of Anxiety Patterns (Imputed Data)", color="white")
        plot_path = os.path.join(output_dir, "hypergraph_imputed.png")
        plt.savefig(plot_path)
        plt.close()
        return f"Hypergraph visualization generated and saved to '{plot_path}'. Representing anxiety patterns."

    except KeyError as ke:
        print(f"KeyError in hypergraph creation: Column '{ke}' not found in DataFrame.")
        return f"Error: Hypergraph generation failed. Column not found: {ke}."
    except Exception as e:
        print(f"Error during hypergraph generation: {e}")
        return "Error: Hypergraph generation failed."

def perform_bootstrap_analysis(data_series, n_resamples=BOOTSTRAP_RESAMPLES):
    """Performs bootstrap analysis to estimate confidence intervals for the mean.

    Args:
        data_series (pd.Series): Data for bootstrap analysis.
        n_resamples (int): Number of bootstrap resamples. Defaults to BOOTSTRAP_RESAMPLES.

    Returns:
        tuple: Confidence interval (low, high) or (None, None) if error occurs.
    """
    try:
        bootstrap_result = bootstrap((data_series,), np.mean, n_resamples=n_resamples, method='percentile', random_state=42)
        confidence_interval = bootstrap_result.confidence_interval
        return confidence_interval
    except Exception as e:
        print(f"Error during bootstrap analysis: {e}")
        return (None, None)

def save_analysis_summary(dataframe, bootstrap_ci, imputation_description, output_dir):
    """Saves summary statistics and imputation details to a text file.

    Args:
        dataframe (pd.DataFrame): DataFrame used for analysis.
        bootstrap_ci (tuple): Confidence interval from bootstrap analysis.
        imputation_description (str): Description of the imputation method used.
        output_dir (str): Directory to save the summary file.

    Returns:
        str: Summary text that was saved, or error message.
    """
    try:
        summary_text = (
            "Descriptive Statistics:\n" + dataframe.describe().to_string() +
            f"\n\nBootstrap Confidence Interval for anxiety_post mean: {bootstrap_ci}" +
            f"\n\nMissing Data Handling Summary:\n{imputation_description}"
        )
        summary_file_path = os.path.join(output_dir, "analysis_summary.txt")
        with open(summary_file_path, "w") as file:
            file.write(summary_text)
        return summary_text
    except Exception as e:
        print(f"Error saving analysis summary: {e}")
        return "Error: Could not save analysis summary."

def generate_insights_report_with_llms(summary_stats_text, shap_analysis_desc, kde_plot_desc, violin_plot_desc, parallel_coords_desc, hypergraph_desc, imputation_desc, output_dir):
    """Generates a comprehensive insights report using placeholder LLM analyses.

    Args:
        summary_stats_text (str): Summary statistics text.
        shap_analysis_desc (str): Description of SHAP analysis.
        kde_plot_desc (str): Description of KDE plot.
        violin_plot_desc (str): Description of violin plots.
        parallel_coords_desc (str): Description of parallel coordinates plots.
        hypergraph_desc (str): Description of hypergraph visualization.
        imputation_desc (str): Description of imputation process.
        output_dir (str): Directory to save the insights report.

    Returns:
        str: Status message indicating success or failure.
    """
    try:
        grok_base_insights = (
            analyze_text_with_llm_placeholder(f"Analyze summary statistics on imputed data:\n{summary_stats_text}", MODEL_GROK_BASE_NAME) + "\n\n" +
            analyze_text_with_llm_placeholder(f"Explain SHAP summary on imputed data: {shap_analysis_desc}", MODEL_GROK_BASE_NAME) + "\n\n" +
            analyze_text_with_llm_placeholder(f"Describe Missing Data Handling approach and implications:\n{imputation_desc}", MODEL_GROK_BASE_NAME)
        )

        claude_sonnet_insights = (
            analyze_text_with_llm_placeholder(f"Interpret KDE plot on imputed data: {kde_plot_desc}", MODEL_CLAUDE_SONNET_NAME) + "\n\n" +
            analyze_text_with_llm_placeholder(f"Interpret Violin plot on imputed data: {violin_plot_desc}", MODEL_CLAUDE_SONNET_NAME) + "\n\n" +
            analyze_text_with_llm_placeholder(f"Interpret Parallel Coordinates Plot on imputed data: {parallel_coords_desc}", MODEL_CLAUDE_SONNET_NAME) + "\n\n" +
            analyze_text_with_llm_placeholder(f"Interpret Hypergraph on imputed data: {hypergraph_desc}", MODEL_CLAUDE_SONNET_NAME)
        )

        grok_enhanced_insights = analyze_text_with_llm_placeholder(
            f"Provide enhanced insights on anxiety intervention effectiveness based on analysis with missing data handling, SHAP, and visualizations on imputed data, focusing on robustness and potential biases mitigated by imputation.",
            MODEL_GROK_ENHANCED_NAME
        )

        combined_insights_report = f"""
    # Combined Insights Report: Anxiety Intervention Analysis with Missing Data Handling (Iterative Imputation)

    ## Grok-base Analysis:
    {grok_base_insights}

    ## Claude 3.7 Sonnet Analysis:
    {claude_sonnet_insights}

    ## Grok-Enhanced Analysis (Missing Data Focused):
    {grok_enhanced_insights}

    ## Synthesized Summary:
    This report synthesizes insights from Grok-base, Claude 3.7 Sonnet, and Grok-Enhanced, focusing on the impact of missing data handling on the anxiety intervention analysis. Iterative imputation using a Bayesian Ridge estimator was employed to address missing values. Grok-base provides a statistical overview and interprets the implications of the iterative imputation approach, noting the importance of pre-anxiety. Claude 3.7 Sonnet details visual patterns and distributions in the imputed dataset, highlighting group differences and shifts in anxiety. Grok-Enhanced, with a missing data focus, delivers nuanced interpretations and actionable recommendations based on the analysis of imputed data, SHAP values, and visualizations, emphasizing the robustness and reduced bias achieved through iterative imputation. The combined expert analyses provide a more reliable and comprehensive assessment of the intervention's effectiveness, addressing potential biases from incomplete data and enabling more accurate causal inferences.
    """

        report_file_path = os.path.join(output_dir, "insights_report.txt")
        with open(report_file_path, "w") as file:
            file.write(combined_insights_report)

        print(f"Insights report generated and saved to: {report_file_path}")
        return "Insights report generated successfully."

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

# --- Main Execution Script ---
if __name__ == "__main__":
    # 1. Prepare Output Directory
    if not create_directory(OUTPUT_PATH):
        exit(1) # Exit with error code

    # 2. Synthetic Dataset with Missing Values
    synthetic_data_csv = """
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,NaN
P005,Group B,5,4
P006,Group B,7,6
P007,Control,3,3
P008,Control,4,NaN
P009,Control,2,2
P010,Control,5,5
"""

    # 3. Load and Validate Data
    anxiety_dataframe = load_dataframe_from_csv_string(synthetic_data_csv)
    if anxiety_dataframe is None:
        exit(1) # Exit if data loading fails

    required_columns = [PARTICIPANT_ID_COLUMN, GROUP_COLUMN, ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN]
    if not validate_dataframe_structure(anxiety_dataframe, required_columns):
        print("Data validation failed. Please check the input data.")
        exit(1) # Exit if validation fails

    # 4. Handle Missing Data using Iterative Imputation
    dataframe_imputed, imputation_description = handle_missing_data(anxiety_dataframe.copy(), imputation_method="iterative")
    if dataframe_imputed is None:
        print(f"Data imputation failed: {imputation_description}")
        exit(1) # Exit if imputation fails

    # 5. One-Hot Encode Group Variable
    original_group_series = dataframe_imputed[GROUP_COLUMN].copy() # Keep original group before encoding
    dataframe_imputed = pd.get_dummies(dataframe_imputed, columns=[GROUP_COLUMN], prefix=GROUP_COLUMN, drop_first=False)
    encoded_group_columns = [col for col in dataframe_imputed.columns if col.startswith(f"{GROUP_COLUMN}_")]
    dataframe_imputed['original_group'] = original_group_series # Add original group back

    # 6. Scale Numerical Data
    columns_to_scale = [ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN] + encoded_group_columns
    dataframe_scaled = scale_dataframe_columns(dataframe_imputed, columns_to_scale)
    if dataframe_scaled is None:
        exit(1) # Exit if scaling fails
    dataframe_for_analysis = dataframe_scaled # Use scaled dataframe for further analysis

    # 7. SHAP Analysis
    shap_feature_cols = encoded_group_columns + [ANXIETY_PRE_COLUMN]
    shap_analysis_info = calculate_shap_and_visualize(
        dataframe_for_analysis.copy(), shap_feature_cols, ANXIETY_POST_COLUMN, OUTPUT_PATH
    )

    # 8. Visualizations
    kde_plot_desc = create_kde_visualization(
        dataframe_for_analysis, ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN, OUTPUT_PATH, NEON_COLORS[:2]
    )
    violin_plot_desc = create_violin_visualizations_by_group(
        dataframe_for_analysis, 'original_group_', ANXIETY_POST_COLUMN, OUTPUT_PATH, NEON_COLORS
    )
    parallel_coords_desc = create_parallel_coordinates_visualization(
        dataframe_for_analysis, 'original_group_', ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN, OUTPUT_PATH, NEON_COLORS
    )
    hypergraph_desc = create_hypergraph_visualization(
        dataframe_for_analysis, ANXIETY_PRE_COLUMN, ANXIETY_POST_COLUMN, OUTPUT_PATH, NEON_COLORS[:2]
    )

    # 9. Bootstrap Analysis
    bootstrap_confidence_interval = perform_bootstrap_analysis(dataframe_for_analysis[ANXIETY_POST_COLUMN], n_resamples=BOOTSTRAP_RESAMPLES)

    # 10. Save Analysis Summary
    summary_stats_text = save_analysis_summary(
        dataframe_for_analysis, bootstrap_confidence_interval, imputation_description, OUTPUT_PATH
    )

    # 11. Generate Insights Report with LLMs
    insights_report_status = generate_insights_report_with_llms(
        summary_stats_text, shap_analysis_info, kde_plot_desc, violin_plot_desc, parallel_coords_desc, hypergraph_desc, imputation_description, OUTPUT_PATH
    )

    print("\n--- Analysis Execution Completed ---")
    print(f"Status of Insights Report Generation: {insights_report_status}")
    print(f"All output files saved in: {OUTPUT_PATH}")