io.py

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

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1l_L5PLrgcvgEnryHe5pFKA47itfa9W4u
"""

!pip install unsloth causal-learn shap transformers plotly accelerate bitsandbytes
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, LabelEncoder
from sklearn.ensemble import RandomForestRegressor
import plotly.express as px
from scipy.stats import bootstrap
import numpy as np
from causallearn.search.ConstraintBased import FCI
from causallearn.utils.GraphUtils import GraphUtils
from transformers import AutoModelForCausalLM, AutoTokenizer
from io import StringIO
from typing import List, Dict, Tuple, Optional, Any
import torch

# --- GPU Check and Unsloth Import ---
try:
    from unsloth import FastLanguageModel
    if not torch.cuda.is_available():
        print("Warning: NVIDIA GPU not found. Unsloth and LLM functionalities will be disabled.")
        USE_UNSLOTH_LLM = False
    else:
        print("NVIDIA GPU found. Unsloth and LLM functionalities will be enabled.")
        USE_UNSLOTH_LLM = True
except ImportError:
    print("Warning: unsloth library not found. LLM functionalities will be disabled.")
    USE_UNSLOTH_LLM = False
except NotImplementedError as e:
    print(f"Warning: {e} LLM functionalities will be disabled.")
    USE_UNSLOTH_LLM = False
except Exception as e:
    print(f"Warning: An unexpected error occurred during unsloth import: {e}. LLM functionalities will be disabled.")
    USE_UNSLOTH_LLM = False


# --- Google Drive Configuration ---
from google.colab import drive
drive.mount('/content/drive')

# --- Configuration ---

# Output Path in Google Drive
OUTPUT_PATH = "/content/drive/MyDrive/output_errors_mistral_grok/"

# DataFrame Column Names
PARTICIPANT_ID_COLUMN: str = "participant_id"
GROUP_COLUMN: str = "group"
ERRORS_COLUMN: str = "errors"
TASK_DIFFICULTY_COLUMN: str = "task_difficulty"
ERROR_TYPE_1_COLUMN: str = "error_type_1"
ERROR_TYPE_2_COLUMN: str = "error_type_2"
ERROR_TYPE_3_COLUMN: str = "error_type_3"
TASK_FEATURE_1_COLUMN: str = "task_feature_1"
TASK_FEATURE_2_COLUMN: str = "task_feature_2"
TASK_FEATURE_3_COLUMN: str = "task_feature_3"

# LLM Models
# We'll use Unsloth for efficient fine-tuning
MODEL_MISTRAL_NAME: str = "unsloth/mistral-7b-v0.3-bnb-4bit"  # Example Unsloth model
MODEL_GROK_NAME: str = "unsloth/mistral-7b-v0.3-bnb-4bit" # Using the same for simplicity, can be different
MODEL_GROK_ENHANCED_NAME: str = "unsloth/mistral-7b-v0.3-bnb-4bit" # Using the same for simplicity, can be different

LINE_WIDTH: float = 2.5
BOOTSTRAP_RESAMPLES: int = 500

# --- Synthetic Dataset ---
SYNTHETIC_DATASET: str = """
participant_id,group,errors,task_difficulty,error_type_1,error_type_2,error_type_3,task_feature_1,task_feature_2,task_feature_3
P001,Group 1,5,7,2,1,0,0.8,0.5,3
P002,Group 1,3,6,1,0,2,0.6,0.3,2
P003,Group 1,6,8,3,2,1,0.9,0.7,4
P004,Group 2,8,9,4,3,0,0.7,0.9,5
P005,Group 2,4,7,0,2,1,0.5,0.4,1
P006,Group 2,7,8,3,1,2,0.8,0.6,3
P007,Group 1,2,5,1,0,0,0.4,0.2,2
P008,Group 2,9,9,4,4,3,0.9,0.8,4
P009,Group 1,4,6,2,1,1,0.7,0.5,3
P010,Group 2,6,7,3,0,2,0.6,0.7,2
P011,Group 3,7,8,2,3,1,0.8,0.9,4
P012,Group 3,5,6,1,2,0,0.7,0.6,3
P013,Group 3,3,5,0,1,2,0.5,0.4,1
P014,Group 3,8,9,4,3,3,0.9,0.8,5
P015,Group 1,6,7,2,2,1,0.8,0.7,3
P016,Group 2,4,6,1,0,2,0.6,0.5,2
P017,Group 1,9,9,4,4,0,0.9,0.9,4
P018,Group 2,2,5,0,1,1,0.4,0.3,1
P019,Group 3,7,8,3,2,2,0.8,0.7,3
P020,Group 1,5,7,2,1,0,0.7,0.6,2
"""

# --- Helper Functions ---

def create_output_directory(path: str) -> None:
    """Creates the output directory if it doesn't exist."""
    os.makedirs(path, exist_ok=True)

def load_data_from_string(csv_string: str) -> pd.DataFrame:
    """Loads data from a CSV string."""
    csv_file = StringIO(csv_string)
    return pd.read_csv(csv_file)

def validate_dataframe(df: pd.DataFrame, required_columns: List[str]) -> bool:
    """Validates the DataFrame: checks for missing columns, data types,
    duplicate IDs, and valid group names.
    """
    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 not in (PARTICIPANT_ID_COLUMN, GROUP_COLUMN):
            if not pd.api.types.is_numeric_dtype(df[col]):
                print(f"Error: Non-numeric data in column: {col}")
                return False

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

    valid_groups = ["Group 1", "Group 2", "Group 3"]
    invalid_groups = df[~df[GROUP_COLUMN].isin(valid_groups)][GROUP_COLUMN].unique()
    if invalid_groups.size > 0:
        print(f"Error: Invalid group names found: {invalid_groups}")
        return False

    return True


def analyze_text_with_llm(text: str, model_name: str, model, tokenizer) -> str:
    """Analyzes the given text using the specified LLM (loaded model)."""
    try:
        inputs = tokenizer(text, return_tensors="pt").to("cuda") # Move inputs to GPU
        outputs = model.generate(**inputs, max_new_tokens=150, do_sample=True)
        return tokenizer.decode(outputs[0], skip_special_tokens=True)
    except Exception as e:
        print(f"Error with LLM {model_name}: {e}")
        return f"Analysis unavailable ({model_name})."

# --- DDQN Implementation (Improved) ---
class DDQNAgent:
    def __init__(self, state_dims: List[int], action_dim: int):
        self.state_dims = state_dims
        self.state_dim = sum(state_dims)  # Total dimension of the flattened state
        self.action_dim = action_dim
        # Initialize Q-network and Target network with random values
        self.q_network = np.random.rand(self.state_dim, action_dim)
        self.target_network = np.random.rand(self.state_dim, action_dim)
        self.update_target_network() # Ensure target network is initialized

    def act(self, state: np.ndarray, epsilon: float = 0.1) -> int:
        """Chooses an action using an epsilon-greedy policy."""
        if np.random.rand() < epsilon:
            return np.random.choice(self.action_dim)  # Explore
        else:
            return np.argmax(self.q_network.T @ state)  # Exploit: Use dot product

    def learn(self, batch: List[Tuple[np.ndarray, int, float, np.ndarray]], gamma: float = 0.99, learning_rate: float = 0.01) -> None:
        """Updates the Q-network using a batch of experiences."""
        for state, action, reward, next_state in batch:
            # Calculate target Q-value
            q_next = self.target_network.T @ next_state
            q_target = reward + gamma * np.max(q_next)

            # Calculate predicted Q-value
            q_predict = self.q_network.T[action] @ state

            # Update Q-network
            self.q_network.T[action] += learning_rate * (q_target - q_predict) * state


    def update_target_network(self) -> None:
        """Copies the weights from the Q-network to the target network."""
        self.target_network = np.copy(self.q_network)

# --- Data Preprocessing ---

def scale_data(df: pd.DataFrame, columns: List[str]) -> pd.DataFrame:
    """Scales the specified columns using MinMaxScaler."""
    scaler = MinMaxScaler()
    df[columns] = scaler.fit_transform(df[columns])
    return df

def preprocess_data(df: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Preprocesses the data by performing one-hot encoding on the 'group' column
    and scaling the numerical columns. Returns both the transformed DataFrame
    and a copy of the original DataFrame.
    """
    df_original = df.copy()  # Store a copy of the original DataFrame
    df = pd.get_dummies(df, columns=[GROUP_COLUMN], prefix=GROUP_COLUMN)
    encoded_group_cols = [col for col in df.columns if col.startswith(f"{GROUP_COLUMN}_")]
    df = scale_data(
        df,
        [
            ERRORS_COLUMN,
            TASK_DIFFICULTY_COLUMN,
            ERROR_TYPE_1_COLUMN,
            ERROR_TYPE_2_COLUMN,
            ERROR_TYPE_3_COLUMN,
            TASK_FEATURE_1_COLUMN,
            TASK_FEATURE_2_COLUMN,
            TASK_FEATURE_3_COLUMN,
        ]
        + encoded_group_cols,
    )
    return df, df_original

# --- Causal Structure Discovery ---

def discover_causal_structure(df: pd.DataFrame, variables: List[str], output_path: str) -> Optional[str]:
    """Discovers the causal structure using FCI and saves the graph."""
    try:
        df_encoded = df.copy()
        for col in df_encoded.select_dtypes(include="object").columns:
            le = LabelEncoder()
            df_encoded[col] = le.fit_transform(df_encoded[col])
        data_fci = df_encoded[variables].to_numpy()
        cg = FCI.fci(data_fci)

        # Check if a graph was actually found
        if cg is None or cg[0] is None or cg[0].graph is None:
            print("Warning: FCI did not find a causal graph.")
            return None

        pdy = GraphUtils.to_pydot(cg[0])  # Access the graph correctly
        pdy.write_png(os.path.join(output_path, "causal_graph.png"))
        edges_info = [
            f"Node {variables[i]}->Node {variables[j]}: Type {edge_type}"
            for i, adj in enumerate(cg[0].graph)
            for j, edge_type in enumerate(adj)
            if edge_type != 0
        ]
        return "\n".join(edges_info)
    except Exception as e:
        print(f"Error in causal discovery: {e}")
        return None

# --- Feature Importance (SHAP) ---

def calculate_shap_values(df: pd.DataFrame, feature_columns: List[str], target_column: str, output_path: str) -> Optional[str]:
    """Calculates and visualizes SHAP values."""
    try:
        df_encoded = df.copy()
        for col in feature_columns:
            if df_encoded[col].dtype == "object":
                le = LabelEncoder()
                df_encoded[col] = le.fit_transform(df_encoded[col])

        model_rf = RandomForestRegressor(random_state=42).fit(
            df_encoded[feature_columns], df_encoded[target_column]
        )
        explainer = shap.TreeExplainer(model_rf)
        shap_values = explainer.shap_values(df_encoded[feature_columns])

        plt.figure(figsize=(10, 8))
        plt.style.use("dark_background")
        shap.summary_plot(
            shap_values, df_encoded[feature_columns], show=False, color_bar=True
        )
        plt.tight_layout()
        plt.savefig(f"{output_path}shap_summary.png")
        plt.close()
        return f"SHAP for {feature_columns} predicting {target_column}"
    except Exception as e:
        print(f"Error calculating SHAP values: {e}")
        return None

# --- Visualization ---

def create_kde_plot(df: pd.DataFrame, column1: str, column2: str, output_path: str, colors: List[str]) -> None:
    """Creates a KDE plot of two columns."""
    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", fontsize=16, color="white")
    plt.legend(facecolor="black", edgecolor="white", labelcolor="white")
    plt.grid(alpha=0.2, linestyle="--")
    plt.tight_layout()
    plt.savefig(f"{output_path}kde_plot.png")
    plt.close()


def create_violin_plot(df: pd.DataFrame, x_column: str, y_column: str, output_path: str, colors: List[str]) -> None:
    """Creates a violin plot."""
    plt.figure(figsize=(10, 6))
    plt.style.use("dark_background")
    sns.violinplot(
        data=df,
        x=x_column,
        y=y_column,
        palette=colors[:2],
        linewidth=LINE_WIDTH,
        hue=x_column,
        legend=False,
    )
    plt.title("Violin Plot", fontsize=16, color="white")
    plt.xlabel(x_column.capitalize(), fontsize=14, color="white")
    plt.ylabel(y_column.capitalize(), fontsize=14, color="white")
    plt.grid(alpha=0.2, linestyle="--")
    plt.tight_layout()
    plt.savefig(f"{output_path}violin_plot.png")
    plt.close()


def create_radar_plot(df: pd.DataFrame, r_column: str, theta_column: str, output_path: str, colors: List[str]) -> None:
    """Creates a radar plot."""
    df_copy = df.copy()
    df_copy[theta_column] = df_copy[theta_column].astype(str)
    fig = px.line_polar(
        df_copy,
        r=r_column,
        theta=theta_column,
        line_close=True,
        color_discrete_sequence=colors,
        template="plotly_dark",
    )
    fig.update_traces(line=dict(width=LINE_WIDTH))
    fig.update_layout(title="Radar Plot")
    try:
        fig.write_image(f"{output_path}radar_plot.png")
    except Exception as e:
        print(f"Plot export failed: {e}")

def create_visualizations(df: pd.DataFrame, df_original: pd.DataFrame, output_path: str, colors: List[str]) -> None:
    """Creates KDE, violin, and radar plots."""
    create_kde_plot(df, ERRORS_COLUMN, TASK_DIFFICULTY_COLUMN, output_path, colors[:2])
    print(f"KDE plot of {ERRORS_COLUMN} and {TASK_DIFFICULTY_COLUMN} created.")
    create_violin_plot(df_original, GROUP_COLUMN, ERRORS_COLUMN, output_path, colors)  # Use original df for group
    print(f"Violin plot: {ERRORS_COLUMN} by {GROUP_COLUMN} created.")
    create_radar_plot(df_original, ERRORS_COLUMN, PARTICIPANT_ID_COLUMN, output_path, colors[:3])  # Use original df for errors
    print(f"Radar plot: {ERRORS_COLUMN} vs {PARTICIPANT_ID_COLUMN} created.")


# --- Hypergraph and Bipartite Graph Conversion ---

def create_feature_nodes(df: pd.DataFrame, shap_values: np.ndarray, feature_names: List[str], threshold_factor: float = 0.5) -> Dict[str, List[str]]:
    """Creates feature nodes for the bipartite graph based on SHAP values."""
    feature_nodes: Dict[str, List[str]] = {}
    mean_shap_values = np.abs(shap_values).mean(axis=0)
    threshold = threshold_factor * np.mean(mean_shap_values)
    for i, feature_name in enumerate(feature_names):
        if mean_shap_values[i] > threshold:
            if df[feature_name].dtype in [np.float64, np.int64]:
                median_val = df[feature_name].median()
                feature_nodes[f"{feature_name}_High"] = (
                    df[df[feature_name] >= median_val][PARTICIPANT_ID_COLUMN].tolist()
                )
                feature_nodes[f"{feature_name}_Low"] = (
                    df[df[feature_name] < median_val][PARTICIPANT_ID_COLUMN].tolist()
                )
            else:
                for unique_val in df[feature_name].unique():
                    feature_nodes[f"{feature_name}_{unique_val}"] = (
                        df[df[feature_name] == unique_val][
                            PARTICIPANT_ID_COLUMN
                        ].tolist()
                    )
    return feature_nodes

def visualize_bipartite_graph(df: pd.DataFrame, feature_nodes: Dict[str, List[str]], output_path: str, colors: List[str]) -> None:
    """Visualizes the bipartite graph."""
    G = nx.Graph()
    participant_ids = df[PARTICIPANT_ID_COLUMN].tolist()
    G.add_nodes_from(participant_ids, bipartite=0)
    feature_node_names = list(feature_nodes.keys())
    G.add_nodes_from(feature_node_names, bipartite=1)
    for feature, participants in feature_nodes.items():
        for participant in participants:
            G.add_edge(participant, feature)
    plt.figure(figsize=(12, 10))
    plt.style.use("dark_background")
    pos = nx.bipartite_layout(G, participant_ids)
    node_colors = [
        colors[0] if node in participant_ids else colors[1] for node in G.nodes()
    ]
    nx.draw(
        G,
        pos,
        with_labels=True,
        node_color=node_colors,
        font_color="white",
        edge_color="gray",
        width=LINE_WIDTH / 2,
        node_size=700,
        font_size=10,
        alpha=0.9,
    )
    plt.title("Bipartite Graph", fontsize=16, color="white")
    plt.tight_layout()
    plt.savefig(f"{output_path}bipartite_graph.png")
    plt.close()


# --- Statistical Analysis ---

def perform_bootstrap(data: pd.Series, statistic: callable, n_resamples: int = BOOTSTRAP_RESAMPLES) -> Tuple[float, float]:
    """Performs bootstrap resampling and returns the confidence interval."""
    result = bootstrap(
        (data,), statistic, n_resamples=n_resamples, method="percentile", random_state=42
    )
    return result.confidence_interval


def save_summary(df: pd.DataFrame, bootstrap_ci: Tuple[float, float], output_path: str) -> str:
    """Calculates and saves summary statistics."""
    summary_text = df.describe().to_string() + f"\nBootstrap CI: {bootstrap_ci}"
    with open(f"{output_path}summary.txt", "w") as f:
        f.write(summary_text)
    return summary_text

def perform_statistical_analysis(df: pd.DataFrame, output_path: str) -> str:
    """Performs statistical analysis and saves the summary."""
    bootstrap_ci = perform_bootstrap(df[ERRORS_COLUMN], np.mean)
    summary_stats_text = save_summary(df, bootstrap_ci, output_path)
    return summary_stats_text

def generate_insights_report(summary_text: str, causal_info: Optional[str], shap_info: Optional[str], kde_desc: str, violin_desc: str, radar_desc: str, bipartite_desc: str, output_path: str, model_mistral=None, tokenizer_mistral=None, model_grok=None, tokenizer_grok=None, model_grok_enhanced=None, tokenizer_grok_enhanced=None) -> None:
    """Generates an insights report using LLMs (using loaded models)."""

    try:
        mistral_insights = "LLM Insights Unavailable (Mistral)."
        grok_insights = "LLM Insights Unavailable (Grok)."
        grok_enhanced_insights = "LLM Insights Unavailable (Grok Enhanced)."

        if USE_UNSLOTH_LLM and model_mistral and tokenizer_mistral:
            mistral_insights = analyze_text_with_llm(
                f"Analyze: {summary_text}\n{causal_info}", MODEL_MISTRAL_NAME, model_mistral, tokenizer_mistral
            )
        if USE_UNSLOTH_LLM and model_grok and tokenizer_grok:
            grok_insights = analyze_text_with_llm(
                f"Analyze: {kde_desc}\n{violin_desc}\n{radar_desc}\n{bipartite_desc}\n{shap_info}",
                MODEL_GROK_NAME, model_grok, tokenizer_grok
            )
        if USE_UNSLOTH_LLM and model_grok_enhanced and tokenizer_grok_enhanced:
            grok_enhanced_insights = analyze_text_with_llm(
                "Synthesize insights.", MODEL_GROK_ENHANCED_NAME, model_grok_enhanced, tokenizer_grok_enhanced
            )

        combined_insights = (
            f"Mistral: {mistral_insights}\nGrok: {grok_insights}\nGrok-Enhanced: {grok_enhanced_insights}"
        )

        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')}")

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



# --- Main Script ---

if __name__ == "__main__":
    create_output_directory(OUTPUT_PATH)

    # Load data
    df = load_data_from_string(SYNTHETIC_DATASET)

    required_columns = [
        PARTICIPANT_ID_COLUMN,
        GROUP_COLUMN,
        ERRORS_COLUMN,
        TASK_DIFFICULTY_COLUMN,
        ERROR_TYPE_1_COLUMN,
        ERROR_TYPE_2_COLUMN,
        ERROR_TYPE_3_COLUMN,
        TASK_FEATURE_1_COLUMN,
        TASK_FEATURE_2_COLUMN,
        TASK_FEATURE_3_COLUMN,
    ]
    if not validate_dataframe(df, required_columns):
        exit(1)  # Exit with an error code

    # Data Preprocessing
    df, df_original = preprocess_data(df)

    # Causal Structure Discovery
    causal_variables = [
        col for col in df.columns if col.startswith(f"{GROUP_COLUMN}_")
    ] + [
        ERRORS_COLUMN,
        TASK_DIFFICULTY_COLUMN,
        ERROR_TYPE_1_COLUMN,
        ERROR_TYPE_2_COLUMN,
        ERROR_TYPE_3_COLUMN,
        TASK_FEATURE_1_COLUMN,
        TASK_FEATURE_2_COLUMN,
        TASK_FEATURE_3_COLUMN,
    ]
    causal_edges_info = discover_causal_structure(
        df.copy(), causal_variables, OUTPUT_PATH
    )

    # Feature Importance (SHAP)
    shap_features = [
        col for col in df.columns if col.startswith(f"{GROUP_COLUMN}_")
    ] + [
        TASK_DIFFICULTY_COLUMN,
        ERROR_TYPE_1_COLUMN,
        ERROR_TYPE_2_COLUMN,
        ERROR_TYPE_3_COLUMN,
        TASK_FEATURE_1_COLUMN,
        TASK_FEATURE_2_COLUMN,
        TASK_FEATURE_3_COLUMN,
    ]
    shap_analysis_info = calculate_shap_values(
        df, shap_features, ERRORS_COLUMN, OUTPUT_PATH
    )

    # Visualization
    plt.style.use("dark_background")
    neon_colors = ["#FF00FF", "#00FFFF", "#FFFF00", "#00FF00"]

    create_visualizations(df, df_original, OUTPUT_PATH, neon_colors)

    # Bipartite Graph
    if shap_analysis_info: # Only create if SHAP values were successful
        model_rf = RandomForestRegressor(random_state=42).fit(
            df[shap_features], df[ERRORS_COLUMN]
        )
        explainer = shap.TreeExplainer(model_rf)
        shap_values = explainer.shap_values(df[shap_features])
        feature_nodes = create_feature_nodes(df, shap_values, shap_features)
        visualize_bipartite_graph(df, feature_nodes, OUTPUT_PATH, neon_colors)
        print("Bipartite graph of participants and features created.")


    # Statistical Analysis
    summary_stats_text = perform_statistical_analysis(df, OUTPUT_PATH)


    # Load LLMs using Unsloth if GPU is available
    model_mistral, tokenizer_mistral, model_grok, tokenizer_grok, model_grok_enhanced, tokenizer_grok_enhanced = None, None, None, None, None, None

    if USE_UNSLOTH_LLM:
        max_seq_length = 2048  # Choose a suitable sequence length
        dtype = None  # Auto-detect dtype
        load_in_4bit = True  # Use 4-bit quantization

        # Mistral
        try:
            model_mistral, tokenizer_mistral = FastLanguageModel.from_pretrained(
                model_name=MODEL_MISTRAL_NAME,
                max_seq_length=max_seq_length,
                dtype=dtype,
                load_in_4bit=load_in_4bit,
            )
        except Exception as e:
            print(f"Error loading Mistral model: {e}")

        # Grok (using Mistral as a stand-in, replace if you have a Grok model)
        try:
            model_grok, tokenizer_grok = FastLanguageModel.from_pretrained(
                model_name=MODEL_GROK_NAME,
                max_seq_length=max_seq_length,
                dtype=dtype,
                load_in_4bit=load_in_4bit,
            )
        except Exception as e:
            print(f"Error loading Grok model: {e}")

        # Grok-Enhanced (using Mistral as a stand-in)
        try:
            model_grok_enhanced, tokenizer_grok_enhanced = FastLanguageModel.from_pretrained(
                model_name=MODEL_GROK_ENHANCED_NAME,
                max_seq_length=max_seq_length,
                dtype=dtype,
                load_in_4bit=load_in_4bit,
            )
        except Exception as e:
            print(f"Error loading Grok Enhanced model: {e}")


    # Generate Insights Report
    generate_insights_report(
        summary_stats_text,
        causal_edges_info,
        shap_analysis_info,
        f"KDE plot of {ERRORS_COLUMN} and {TASK_DIFFICULTY_COLUMN}",
        f"Violin plot: {ERRORS_COLUMN} by {GROUP_COLUMN}",
        f"Radar plot: {ERRORS_COLUMN} vs {PARTICIPANT_ID_COLUMN}",
        "Bipartite graph of participants and features",
        OUTPUT_PATH,
        model_mistral, tokenizer_mistral,
        model_grok, tokenizer_grok,
        model_grok_enhanced, tokenizer_grok_enhanced
    )


    # DDQN Agent Training
    state_dims = [1] * (len(shap_features))  # Each feature is a dimension
    action_dim = 2  # Keep it simple for demonstration
    ddqn_agent = DDQNAgent(state_dims, action_dim)

    for epoch in range(10):
        batch = []
        for _, row in df.iterrows():
            # Create state as a numpy array
            state = np.array([row[col] for col in shap_features])
            action = ddqn_agent.act(state)
            reward = -row[ERRORS_COLUMN]  # Negative of errors
            next_state = state # In this simplified example, next_state = state
            batch.append((state, action, reward, next_state))
        ddqn_agent.learn(batch)
        ddqn_agent.update_target_network()  # Update target network each epoch
    print("DDQN Agent Training Completed")

    # Save Processed Data
    df.to_csv(f"{OUTPUT_PATH}output_data.csv", index=False)
    print("Execution completed successfully.")

    # --- Unsloth Fine-tuning and QLoRA (Example) ---
    # This section demonstrates *how* you would fine-tune.  It's not a complete
    # training loop, but shows the core Unsloth steps.  You would integrate this
    # into a proper training setup with your dataset.

    if USE_UNSLOTH_LLM and model_mistral and tokenizer_mistral: # Only proceed if LLMs and GPU are available
        print("\n--- Fine-tuning section starting ---")
        #  Prepare the model for QLoRA
        try:
            model_mistral = FastLanguageModel.get_peft_model(
                model_mistral,
                r=16,  # LoRA rank
                target_modules=[
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "o_proj",
                    "gate_proj",
                    "up_proj",
                    "down_proj",
                ],
                lora_alpha=16,
                lora_dropout=0,  # Optimized
                bias="none",  # Optimized
                use_gradient_checkpointing="unsloth",  # Optimized
                random_state=3407,
                use_rslora=False,
            )

            #  Create a dummy dataset for demonstration.  Replace with your actual data.
            from datasets import Dataset
            dataset = Dataset.from_dict({
                "text": [
                    "Group 1 has high error rates.",
                    "Task difficulty correlates with errors.",
                    "Group 2 shows fewer errors.",
                    "Error type 1 is common in Group 3."
                ] * 5  # Repeat for more data
            })


            #  Set up the SFTTrainer (from trl)
            from trl import SFTTrainer
            from transformers import TrainingArguments

            trainer = SFTTrainer(
                model=model_mistral,
                tokenizer=tokenizer_mistral,
                train_dataset=dataset,
                dataset_text_field="text",
                max_seq_length=max_seq_length,
                packing=False,  # Can speed up training for short sequences
                args=TrainingArguments(
                    per_device_train_batch_size=2,
                    gradient_accumulation_steps=4,
                    warmup_steps=5,
                    max_steps=10,  # Keep it short for demonstration
                    learning_rate=2e-4,
                    fp16=not torch.cuda.is_bf16_supported(),
                    bf16=torch.cuda.is_bf16_supported(),
                    logging_steps=1,
                    optim="adamw_8bit",
                    weight_decay=0.01,
                    lr_scheduler_type="linear",
                    seed=3407,
                    output_dir=OUTPUT_PATH,
                    report_to="none",  # Disable reporting for this example
                ),
            )

            #  Train (this is a very short example training run)
            trainer.train()

            # Save model and tokenizer
            trainer.save_model(f"{OUTPUT_PATH}fine_tuned_model")
            tokenizer_mistral.save_pretrained(f"{OUTPUT_PATH}fine_tuned_model")
            print(f"Fine-tuned model saved to {OUTPUT_PATH}fine_tuned_model")

        except Exception as e:
            print(f"Error during fine-tuning: {e}")
    else:
        print("\n--- Fine-tuning section skipped due to missing GPU or LLM setup ---")


    # --- GGUF Quantization (using llama.cpp) ---
    # This part requires llama.cpp to be installed and accessible.
    # We'll use a subprocess call for demonstration.  In a real setup,
    # you might use the llama.cpp Python bindings.

    if USE_UNSLOTH_LLM and model_mistral and tokenizer_mistral: # Only proceed if LLMs and GPU are available
        print("\n--- GGUF Quantization section starting ---")
        # 1. Convert the fine-tuned model to HF format (if needed)
        #    Unsloth saves in a compatible format, so this step might be skippable.
        #    But for completeness, here's how you'd generally do it:
        # from transformers import AutoModelForCausalLM
        # model = AutoModelForCausalLM.from_pretrained(f"{OUTPUT_PATH}fine_tuned_model")
        # model.save_pretrained(f"{OUTPUT_PATH}fine_tuned_hf")

        # 2. Convert to GGUF
        import subprocess

        # The path to your llama.cpp 'convert' script.  ADJUST THIS!
        llama_cpp_convert_path = "/path/to/your/llama.cpp/convert.py"  #  **IMPORTANT:  Replace with your actual path!**
        # The path to the fine-tuned HF model (adjust if you did the conversion)
        hf_model_path = f"{OUTPUT_PATH}fine_tuned_model"
        # Output GGUF file path
        gguf_output_path = f"{OUTPUT_PATH}fine_tuned_model.gguf"

        # Construct the conversion command.  We're using Q4_0 quantization.
        convert_command = [
            "python3",
            llama_cpp_convert_path,
            hf_model_path,
            "--outfile",
            gguf_output_path,
            "--outtype",
            "q4_0",  # Quantization type
        ]

        # Run the conversion.  This might take a while!
        try:
            subprocess.run(convert_command, check=True, capture_output=True, text=True)
            print(f"GGUF model saved to: {gguf_output_path}")
        except subprocess.CalledProcessError as e:
            print(f"Error during GGUF conversion:\n{e.stderr}")
        except FileNotFoundError:
            print(f"Error: Could not find the llama.cpp convert script at {llama_cpp_convert_path}.  Please install llama.cpp and update the path.")
        except Exception as e:
            print(f"An unexpected error occurred during GGUF conversion: {e}")
    else:
        print("\n--- GGUF Quantization section skipped due to missing GPU or LLM setup ---")