processor.py

#!/usr/bin/env python3
"""
Execute Processor module for GNN Processing Pipeline.

This module provides execute processing capabilities for rendered implementations.
"""

import json
import logging
import os
import subprocess  # nosec B404 -- subprocess calls with controlled/trusted input
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from utils.logging.logging_utils import (
    log_step_error,
    log_step_start,
    log_step_success,
    log_step_warning,
)

try:
    from utils.logging.logging_utils import PipelineLogger
except ImportError:
    PipelineLogger = None

logger = logging.getLogger(__name__)



def check_julia_dependencies(verbose: bool, log: Optional[logging.Logger] = None) -> bool:
    """Check if required Julia packages are available.

    Args:
        verbose: Enable verbose logging.
        log: Optional logger instance; defaults to module logger if not provided.

    Returns:
        True if dependencies ok, False otherwise.
    """
    if log is None:
        log = logger
    try:
        # check basic julia availability
        subprocess.run(['julia', '--version'], capture_output=True, check=True, timeout=10)  # nosec B607 B603 -- subprocess calls with controlled/trusted input

        # Check for key packages
        check_script = 'using Pkg; Pkg.status(["RxInfer", "ActiveInference", "GraphPPL"])'
        result = subprocess.run(  # nosec B607 B603 -- subprocess calls with controlled/trusted input
            ['julia', '-e', check_script],
            capture_output=True,
            text=True,
            timeout=30
        )

        if result.returncode != 0:
            if verbose:
                log.warning(f"Julia package check failed: {result.stderr}")
            return False

        return True
    except (subprocess.CalledProcessError, FileNotFoundError):
        return False


def determine_script_framework(script_path: Path, render_output_dir: Path, framework_dirs: Dict[str, str]) -> str:
    """
    Determine the framework for a script based on its directory path.

    Args:
        script_path: Path to the script
        render_output_dir: Base render output directory
        framework_dirs: Mapping of directory names to framework names

    Returns:
        Framework name or 'unknown'
    """
    try:
        # Get relative path from render output directory
        relative_path = script_path.relative_to(render_output_dir)

        # Check each part of the path for framework indicators
        for part in relative_path.parts:
            # Check if this part matches a known framework directory
            if part.lower() in framework_dirs:
                return framework_dirs[part.lower()]

            # Check for framework names in the directory name
            for framework_name in framework_dirs.values():
                if framework_name.lower() in part.lower():
                    return framework_name

        # Default recovery
        return "unknown"

    except Exception as e:
        logging.getLogger(__name__).debug(f"Error determining framework for script: {e}")
        return "unknown"


# Phase 2.3: framework-availability helpers moved to utils.framework_availability
# so execute and render stay in sync. The import-check dict and predicate are
# re-exported here via thin aliases to preserve any external callers that
# previously imported them from execute.processor.
from utils.framework_availability import (  # noqa: E402
    FRAMEWORK_IMPORT_CHECK as _FRAMEWORK_IMPORT_CHECK,
    is_framework_available as _is_framework_available_by_name,
)


def _is_python_framework_dependency_available(framework: str, executor: str, logger) -> bool:
    """Return True if the framework's required Python module is importable.

    Delegates to ``utils.framework_availability.is_framework_available``, passing
    ``executor`` so the check targets the subprocess-invoked interpreter rather
    than the caller's. Preserves the pre-Phase-2.3 call-site signature.
    """
    return _is_framework_available_by_name(framework, executor=executor, logger=logger)


def _make_skipped_result(script_info: Dict[str, Any], framework: str, model_name: str, executor: str, logger) -> Dict[str, Any]:
    """Build an execution result dict for a script skipped due to missing dependency."""
    module_name, install_hint = _FRAMEWORK_IMPORT_CHECK.get(framework, ("", ""))
    reason = f"Dependency not installed: {module_name}" if module_name else "Dependency not installed"
    if install_hint and not logger.isEnabledFor(logging.DEBUG):
        logger.info(f"Skipping {script_info['name']} ({framework}): {module_name} not installed. Install with: {install_hint}")
    return {
        "script_path": str(script_info["path"]),
        "script_name": script_info["name"],
        "framework": framework,
        "model_name": model_name,
        "executor": executor,
        "success": False,
        "skipped": True,
        "return_code": None,
        "stdout": "",
        "stderr": "",
        "execution_time": 0,
        "timestamp": datetime.now().isoformat(),
        "error": reason,
        "error_type": "DependencyNotInstalled",
    }


def parse_frameworks_parameter(frameworks: str, logger) -> List[str]:
    """
    Parse the frameworks parameter into a list of framework names.

    Args:
        frameworks: Comma-separated string of framework names or preset
        logger: Logger instance

    Returns:
        List of framework names to include
    """
    if not frameworks or frameworks.lower() == "all":
        return ["pymdp", "jax", "discopy", "rxinfer", "activeinference_jl", "pytorch", "numpyro", "bnlearn"]

    if frameworks.lower() == "lite":
        return ["pymdp", "jax", "discopy", "bnlearn"]

    # Parse comma-separated list
    framework_list = [f.strip() for f in frameworks.split(",")]
    valid_frameworks = ["pymdp", "jax", "discopy", "rxinfer", "activeinference_jl", "pytorch", "numpyro", "bnlearn"]

    # Filter out invalid frameworks
    valid_list = [f for f in framework_list if f in valid_frameworks]

    if len(valid_list) != len(framework_list):
        invalid = [f for f in framework_list if f not in valid_frameworks]
        logger.warning(f"Invalid frameworks specified: {invalid}. Valid options: {valid_frameworks}")

    return valid_list if valid_list else ["pymdp"]  # Default to pymdp if nothing valid


def _resolve_render_output_dir(target_dir: Path, kwargs: dict) -> Optional[Path]:
    """Resolve the render output directory from kwargs and filesystem heuristics.

    Resolution priority:
    1. Explicit ``--render-output-dir`` kwarg.
    2. target_dir itself if it looks like a render output directory.
    3. Common pipeline and test output locations (searched in order).

    Returns the first existing, non-empty directory found, or None.
    """
    # Priority 1: explicit kwarg
    if kwargs.get('render_output_dir'):
        return Path(kwargs['render_output_dir'])

    # Priority 2: target_dir is already the render output
    if "11_render_output" in str(target_dir) or target_dir.name == "11_render_output":
        return target_dir

    # Priority 3: search common locations
    candidates: List[Path] = [
        target_dir.parent / "output" / "11_render_output",
        target_dir / "11_render_output",
        Path("output/test_render/11_render_output/11_render_output"),
        Path("output/test_render_improved/11_render_output/11_render_output"),
        *list(Path("output").glob("*/11_render_output/11_render_output")),
        *list(Path("output").glob("**/11_render_output")),
    ]
    for candidate in candidates:
        if candidate.exists() and any(candidate.rglob("*")):
            return candidate
    return None


def process_execute(
    target_dir: Path,
    output_dir: Path,
    verbose: bool = False,
    frameworks: str = "all",
    **kwargs: Any,
) -> Union[bool, int]:
    """
    Execute rendered implementations from 11_render_output directory.
    
    This function searches for executable scripts generated by 11_render.py
    and executes them using subprocess, capturing their outputs and results.
    
    Args:
        target_dir: Directory containing rendered executable scripts (typically 11_render_output)
        output_dir: Directory to save execution results
        verbose: Enable verbose output
        **kwargs: Additional arguments
        
    Returns:
        True if processing successful, False otherwise
    """
    logger = logging.getLogger("execute")

    try:
        log_step_start(logger, "Processing execute - searching for rendered implementations")

        # Phase 1.3: validate frameworks arg before parsing. Rejects non-string
        # input and fully-unknown framework lists early with a clear error.
        try:
            from utils.validation_schemas import validate_frameworks_arg
            frameworks = validate_frameworks_arg(frameworks, context="process_execute")
        except ValueError as _verr:
            log_step_error(logger, f"Invalid frameworks argument: {_verr}")
            return False

        # Parse frameworks parameter
        requested_frameworks = parse_frameworks_parameter(frameworks, logger)
        logger.info(f"Requested frameworks: {requested_frameworks}")

        results_dir = output_dir
        results_dir.mkdir(parents=True, exist_ok=True)

        # Initialize execution results
        execution_results = {
            "timestamp": datetime.now().isoformat(),
            "target_directory": str(target_dir),
            "output_directory": str(output_dir),
            "total_scripts_found": 0,
            "successful_executions": 0,
            "failed_executions": 0,
            "skipped_executions": 0,
            "execution_details": [],
            "framework_status": {},
            "success": True
        }

        # Look for rendered implementations from render output
        render_output_dir = _resolve_render_output_dir(target_dir, kwargs)
        if render_output_dir is not None and render_output_dir != target_dir:
            logger.info(f"Found render output directory: {render_output_dir}")

        if verbose:
            logger.info(f"Searching for executable scripts in: {render_output_dir}")

        if not render_output_dir or not render_output_dir.exists():
            log_step_warning(logger, f"Render output directory not found: {render_output_dir}")
            execution_results["success"] = True  # Not a hard error
            execution_results["skipped_reason"] = "no_render_output"
            execution_results["message"] = "No rendered implementations found"
        else:
            # Find executable scripts, filtered by requested frameworks
            executable_scripts = find_executable_scripts(render_output_dir, verbose, logger, requested_frameworks)
            execution_results["total_scripts_found"] = len(executable_scripts)
            execution_results["requested_frameworks"] = requested_frameworks

            if not executable_scripts:
                log_step_warning(logger, "No executable scripts found in render output")
                execution_results["message"] = "No executable scripts found"
                execution_results["success"] = True
                execution_results["skipped_reason"] = "no_executable_scripts"
            else:
                logger.info(f"Found {len(executable_scripts)} executable scripts to run")

                # Extract args
                timeout = kwargs.get('timeout', 300)
                is_distributed = kwargs.get('distributed', False)
                details = []

                if is_distributed:
                    from .distributed import Dispatcher
                    backend = kwargs.get('backend', 'ray')
                    dispatcher = Dispatcher(backend=backend)
                    
                    def ray_script_runner(info, **kws):
                        """Execute a rendered simulation script using Ray for distributed processing."""
                        # Re-instantiate logger to avoid pickle issues
                        import logging
                        local_logger = logging.getLogger("execute.worker")
                        local_logger.setLevel(logging.INFO)
                        return execute_single_script(info, kws["results_dir"], kws["verbose"], local_logger, kws["timeout"])
                        
                    details = dispatcher.run_scripts_parallel(
                        executable_scripts, 
                        ray_script_runner, 
                        results_dir=results_dir, 
                        verbose=verbose, 
                        timeout=timeout
                    )
                else:
                    # Execute each script sequentially (execute_single_script skips when optional dep missing)
                    for script_info in executable_scripts:
                        exec_result = execute_single_script(script_info, results_dir, verbose, logger, timeout)
                        exec_result.setdefault("skipped", False)
                        details.append(exec_result)

                # Update aggregated results
                for exec_result in details:
                    execution_results["execution_details"].append(exec_result)

                    # Update framework status
                    framework = exec_result.get("framework", "unknown")
                    if framework not in execution_results["framework_status"]:
                        execution_results["framework_status"][framework] = {"status": "unknown", "executions": 0}

                    execution_results["framework_status"][framework]["executions"] += 1

                    if exec_result.get("skipped"):
                        execution_results["skipped_executions"] = execution_results.get("skipped_executions", 0) + 1
                        execution_results["framework_status"][framework]["status"] = "skipped"
                        if "error" in exec_result:
                            execution_results["framework_status"][framework]["error"] = exec_result["error"]
                    elif exec_result["success"]:
                        execution_results["successful_executions"] += 1
                        execution_results["framework_status"][framework]["status"] = "success"
                    else:
                        execution_results["failed_executions"] += 1
                        execution_results["framework_status"][framework]["status"] = "failed"
                        if "error" in exec_result:
                            execution_results["framework_status"][framework]["error"] = exec_result["error"]

        # Populate backward-compatible keys before saving
        total_found = execution_results["total_scripts_found"]
        successful = execution_results["successful_executions"]
        failed = execution_results["failed_executions"]
        skipped = execution_results.get("skipped_executions", 0)
        attempted = total_found - skipped
        execution_results["total_scripts"] = total_found
        execution_results["success_rate"] = (
            round(successful / attempted * 100, 2) if attempted > 0 else 100.0
        )

        # Save detailed results to summaries subfolder
        summaries_dir = results_dir / "summaries"
        summaries_dir.mkdir(parents=True, exist_ok=True)
        results_file = summaries_dir / "execution_summary.json"
        with open(results_file, 'w') as f:
            json.dump(execution_results, f, indent=2, default=str)

        # Generate execution report
        generate_execution_report(execution_results, results_dir, logger)

        # Determine overall success: only count real failures (not skipped) toward critical threshold
        total_scripts = execution_results["total_scripts_found"]
        failed_scripts = execution_results["failed_executions"]
        skipped_scripts = execution_results.get("skipped_executions", 0)
        attempted_scripts = total_scripts - skipped_scripts

        if total_scripts == 0:
            log_step_warning(logger, "No executable scripts found to run")
            # Exit-code 2: step completed without doing work. Distinguishes
            # "nothing to do" from "did work successfully".
            return 2
        elif failed_scripts == 0:
            if skipped_scripts:
                log_step_success(logger, f"Execute completed: {execution_results['successful_executions']} succeeded, {skipped_scripts} skipped (dependency not installed)")
            else:
                log_step_success(logger, "Execute processing completed successfully")
        elif attempted_scripts > 0 and failed_scripts < attempted_scripts * 0.5:
            log_step_warning(logger, f"Execute completed with {failed_scripts}/{attempted_scripts} failures (partial success)" + (f", {skipped_scripts} skipped" if skipped_scripts else ""))
        elif attempted_scripts > 0:
            log_step_error(logger, f"Execute completed with {failed_scripts}/{attempted_scripts} failures (critical)" + (f", {skipped_scripts} skipped" if skipped_scripts else ""))

        # Return True if we found and attempted to run scripts (even if some failed)
        # Return False only if there was a critical error preventing execution
        return True

    except Exception as e:
        log_step_error(logger, f"Execute processing failed: {e}")
        return False

def find_executable_scripts(render_output_dir: Path, verbose: bool, logger, requested_frameworks: List[str]) -> List[Dict[str, Any]]:
    """
    Find executable scripts in the render output directory.

    Searches for Python (.py) and Julia (.jl) scripts in the render output
    directory structure. Scripts are filtered by the requested frameworks
    and excluded if they match common non-executable patterns (test files,
    __init__.py, etc.).

    Args:
        render_output_dir: Directory containing rendered scripts from Step 11.
        verbose: Enable verbose logging of discovered scripts.
        logger: Logger instance for output messages.
        requested_frameworks: List of framework names to include (e.g.,
            ["pymdp", "jax", "discopy"]). Scripts from other frameworks
            will be skipped.

    Returns:
        List of dictionaries, each containing:
            - path: Path to the script file
            - name: Script filename
            - framework: Detected framework name
            - executor: Command to execute the script (python/julia)
            - relative_path: Path relative to render_output_dir
            - size_bytes: File size in bytes
    """
    executable_scripts = []

    # Define supported script types and their executors
    script_types = {
        '*.py': {'executor': sys.executable, 'framework': 'python'},
        '*.jl': {'executor': 'julia', 'framework': 'julia'},
    }

    # Map framework directories to framework names
    framework_dirs = {
        'pymdp': 'pymdp',
        'jax': 'jax',
        'discopy': 'discopy',
        'rxinfer': 'rxinfer',
        'activeinference_jl': 'activeinference_jl',
        'activeinference.jl': 'activeinference_jl',
        'pytorch': 'pytorch',
        'numpyro': 'numpyro',
        'bnlearn': 'bnlearn',
    }

    for pattern, config in script_types.items():
        scripts = list(render_output_dir.rglob(pattern))

        for script_path in scripts:
            # Skip test files and other non-executable scripts
            if any(skip in script_path.name.lower() for skip in ['test_', '__', 'readme']):
                continue

            # Determine framework from directory path
            framework = determine_script_framework(script_path, render_output_dir, framework_dirs)

            # Filter by requested frameworks
            if framework not in requested_frameworks:
                if verbose:
                    logger.debug(f"Skipping {framework} script: {script_path.name} (not in requested frameworks)")
                continue

            # Check if script is executable or can be made executable
            script_info = {
                'path': script_path,
                'name': script_path.name,
                'framework': framework,
                'executor': config['executor'],
                'relative_path': script_path.relative_to(render_output_dir),
                'size_bytes': script_path.stat().st_size if script_path.exists() else 0
            }

            executable_scripts.append(script_info)

            if verbose:
                logger.info(f"Found {config['framework']} script: {script_info['relative_path']}")

    return executable_scripts


def execute_single_script(script_info: Dict[str, Any], results_dir: Path, verbose: bool, logger, timeout: int = 300) -> Dict[str, Any]:
    """
    Execute a single script using subprocess.
    
    Args:
        script_info: Dictionary containing script information
        results_dir: Directory to save execution results (will create implementation-specific subfolders)
        verbose: Enable verbose logging
        logger: Logger instance
        
    Returns:
        Dictionary with execution results
    """
    script_path = script_info['path']
    executor = script_info['executor']

    # Extract model name and framework from script path for organization
    # Expected path: .../11_render_output/model_name/framework/script.ext
    path_parts = script_path.parts
    if len(path_parts) >= 3:
        model_name = path_parts[-3]  # e.g., 'actinf_pomdp_agent'
        framework = path_parts[-2]   # e.g., 'pymdp'
    else:
        model_name = 'unknown_model'
        framework = script_info['framework']

    # Pre-flight skip: do not run Python frameworks when optional dependency is missing
    if executor == sys.executable and not _is_python_framework_dependency_available(framework, executor, logger):
        return _make_skipped_result(script_info, framework, model_name, executor, logger)

    # Prepare execution result
    exec_result = {
        'script_path': str(script_path),
        'script_name': script_info['name'],
        'framework': framework,
        'model_name': model_name,
        'executor': executor,
        'success': False,
        'return_code': None,
        'stdout': '',
        'stderr': '',
        'execution_time': 0,
        'timestamp': datetime.now().isoformat()
    }

    try:

        if verbose:
            logger.info(f"Executing {script_info['framework']} script: {script_info['name']}")

        start_time = datetime.now()

        # Check if the executor is available
        try:
            # For Python scripts, check if Python is available (most are Python scripts)
            if executor in ['python', 'python3']:
                subprocess.run([executor, '--version'],  # nosec B603 -- subprocess calls with controlled/trusted input
                             capture_output=True,
                             text=True,
                             timeout=5,
                             check=True)

                # For PyMDP, specifically check if it's importable
                if framework == "pymdp":
                    try:
                        import_check = subprocess.run(  # nosec B603 -- subprocess calls with controlled/trusted input
                            [executor, '-c', 'import pymdp; print("ok")'],
                            capture_output=True,
                            text=True,
                            timeout=5
                        )
                        if import_check.returncode != 0:
                            logger.warning(f"PyMDP package appears missing or broken: {import_check.stderr}")
                            exec_result['error'] = f"PyMDP dependency missing: {import_check.stderr}"
                            # Continue anyway as it might be a local import, but log warning
                    except Exception as e:
                        logger.debug(f"Error checking PyMDP importability: {e}")

            # For Julia scripts, check availability and dependencies
            elif executor == 'julia':
                if not check_julia_dependencies(verbose, logger):
                    logger.warning("Julia dependencies (RxInfer/ActiveInference) may be missing")
                    # Don't fail here, let the script try to run, but log warning

                subprocess.run([executor, '--version'],  # nosec B603 -- subprocess calls with controlled/trusted input
                             capture_output=True,
                             text=True,
                             timeout=5,
                             check=True)
            # For other executors, try a basic check
            else:
                subprocess.run([executor, '--version'],  # nosec B603 -- subprocess calls with controlled/trusted input
                             capture_output=True,
                             text=True,
                             timeout=5,
                             check=True)
        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired) as e:
            exec_result['error'] = f"Executor '{executor}' is not available or not working: {e}"
            logger.warning(f"Executor '{executor}' is not available - skipping {script_info['name']}")
            return exec_result

        # Execute the script with improved error handling
        script_name = script_path.name
        result = None

        # Error result class for consistent interface when subprocess fails
        class ErrorResult:
            def __init__(self, returncode: int, stdout: str, stderr: str):
                self.returncode = returncode
                self.stdout = stdout
                self.stderr = stderr

        try:
            # Set environment variables if needed
            env = os.environ.copy()
            if framework == "pymdp":
                env["PYTHONPATH"] = str(script_path.parent) + os.pathsep + env.get("PYTHONPATH", "")
                # Generated PyMDP scripts import src.execute.pymdp; subprocess cwd is often under output/
                _proc = Path(__file__).resolve()
                _repo_root = _proc.parent.parent.parent  # src/execute/processor.py -> src -> repo root
                env["GNN_PROJECT_ROOT"] = str(_repo_root)

            # Direct JAX/NumPyro/PyTorch output into execute output dir for collection/analysis
            impl_output_dir = results_dir / model_name / framework
            sim_data_dir = impl_output_dir / "simulation_data"
            if framework == "jax":
                sim_data_dir.mkdir(parents=True, exist_ok=True)
                env["GNN_OUTPUT_DIR"] = str(sim_data_dir)
            elif framework == "numpyro":
                sim_data_dir.mkdir(parents=True, exist_ok=True)
                env["NUMPYRO_OUTPUT_DIR"] = str(sim_data_dir)
            elif framework == "pytorch":
                sim_data_dir.mkdir(parents=True, exist_ok=True)
                env["PYTORCH_OUTPUT_DIR"] = str(sim_data_dir)

            result = subprocess.run(  # nosec B603 -- subprocess calls with controlled/trusted input
                [executor, script_name],
                capture_output=True,
                text=True,
                timeout=timeout,
                cwd=script_path.parent,  # Run in the script's directory
                env=env
            )

            end_time = datetime.now()
            exec_result['execution_time'] = (end_time - start_time).total_seconds()
            exec_result['return_code'] = result.returncode
            exec_result['stdout'] = result.stdout
            exec_result['stderr'] = result.stderr

            if result.returncode == 0:
                exec_result['success'] = True
                logger.info(f"✅ Successfully executed {script_info['name']}")
                if verbose and result.stdout:
                    logger.info(f"Script output: {result.stdout[:200]}...")  # Show first 200 chars
            else:
                exec_result['error'] = f"Script failed with return code {result.returncode}"

                # Analyze stderr for common errors
                if "ModuleNotFoundError" in result.stderr:
                    exec_result['error_type'] = "DependencyError"
                    logger.error(f"Missing dependency in {script_info['name']}: {result.stderr.splitlines()[-1]}")
                elif "SyntaxError" in result.stderr:
                    exec_result['error_type'] = "SyntaxError"
                    logger.error(f"Syntax error in {script_info['name']}")
                else:
                    exec_result['error_type'] = "RuntimeError"

                logger.warning(f"⚠️ Script {script_info['name']} failed with return code {result.returncode}")
                if result.stderr:
                    logger.warning(f"Error output: {result.stderr[:500]}...")  # Show first 500 chars

        except subprocess.TimeoutExpired:
            end_time = datetime.now()
            exec_result['execution_time'] = (end_time - start_time).total_seconds()
            exec_result['error'] = f"Script execution timed out after {timeout} seconds"
            exec_result['return_code'] = -1
            exec_result['stdout'] = ""
            exec_result['stderr'] = "Timeout"
            logger.warning(f"⏰ Script {script_info['name']} timed out after {timeout} seconds")
            result = ErrorResult(-1, "", "Timeout")

        except Exception as e:
            end_time = datetime.now()
            exec_result['execution_time'] = (end_time - start_time).total_seconds()
            exec_result['error'] = f"Script execution failed: {e}"
            exec_result['return_code'] = -2
            exec_result['stdout'] = ""
            exec_result['stderr'] = str(e)
            logger.warning(f"❌ Script {script_info['name']} execution failed: {e}")
            result = ErrorResult(-2, "", str(e))

        # Ensure result is defined before using it
        if result is None:
            result = ErrorResult(-3, "", "Unknown error")

        # Save individual script output in implementation-specific subdirectory
        # Create the implementation-specific directory structure
        impl_specific_dir = results_dir / model_name / framework / "execution_logs"
        impl_specific_dir.mkdir(parents=True, exist_ok=True)

        # Note: Framework-specific subdirectories (visualizations, simulation_data, etc.)
        # are created on-demand by collect_execution_outputs() only when actual content
        # is copied to them, avoiding empty folder creation.


        # Extract simulation data from stdout/stderr
        simulation_data = _extract_simulation_data(result.stdout, result.stderr, framework, logger)
        exec_result['simulation_data'] = simulation_data

        # Save structured execution results in JSON format
        structured_result = {
            "framework": framework,
            "model_name": model_name,
            "script_name": script_info['name'],
            "script_path": str(script_path),
            "success": exec_result['success'],
            "return_code": exec_result.get('return_code'),
            "execution_time": exec_result.get('execution_time', 0),
            "timestamp": exec_result['timestamp'],
            "simulation_data": simulation_data,
            "execution_metadata": {
                "executor": executor,
                "stdout_length": len(result.stdout),
                "stderr_length": len(result.stderr),
                "output_directory": str(impl_specific_dir.parent)
            }
        }

        # Save structured JSON result
        json_output_file = impl_specific_dir / f"{script_info['name']}_results.json"
        with open(json_output_file, 'w') as f:
            json.dump(structured_result, f, indent=2, default=str)

        exec_result['structured_result_file'] = str(json_output_file)

        # Also save human-readable log
        output_file = impl_specific_dir / f"{script_info['name']}_execution.log"
        with open(output_file, 'w') as f:
            f.write(f"Execution Results for {script_info['name']}\n")
            f.write(f"Timestamp: {exec_result['timestamp']}\n")
            f.write(f"Return Code: {result.returncode}\n")
            f.write(f"Execution Time: {exec_result['execution_time']:.2f} seconds\n")
            f.write(f"Model: {model_name}\n")
            f.write(f"Framework: {framework}\n")
            f.write(f"Output Directory: {impl_specific_dir.parent}\n\n")
            f.write("STDOUT:\n")
            f.write(result.stdout)
            f.write("\n\nSTDERR:\n")
            f.write(result.stderr)

        exec_result['output_file'] = str(output_file)
        exec_result['implementation_directory'] = str(impl_specific_dir.parent)

        # Collect execution outputs (visualizations, simulation data, traces)
        if exec_result['success']:
            try:
                logger.info(f"Collecting execution outputs for {framework} script {script_info['name']}")
                collected_outputs = collect_execution_outputs(
                    script_path,
                    impl_specific_dir.parent,
                    framework,
                    logger
                )
                exec_result['collected_outputs'] = collected_outputs

                # Update structured result with collected file paths
                structured_result['collected_outputs'] = collected_outputs

                # Re-save structured result with collected outputs
                with open(json_output_file, 'w') as f:
                    json.dump(structured_result, f, indent=2, default=str)
                logger.debug("Updated results JSON with collected outputs")

                # Enhance simulation data extraction from collected files
                if collected_outputs:
                    logger.info(f"Extracting simulation data from collected files for {framework}")
                    enhanced_data = _extract_simulation_data_from_files(
                        impl_specific_dir.parent,
                        framework,
                        logger
                    )
                    if enhanced_data:
                        logger.info(f"Extracted {len(enhanced_data)} data fields from files")
                        simulation_data.update(enhanced_data)
                        exec_result['simulation_data'] = simulation_data
                        structured_result['simulation_data'] = simulation_data

                        # Re-save again with enhanced data
                        with open(json_output_file, 'w') as f:
                            json.dump(structured_result, f, indent=2, default=str)
                        logger.debug("Updated results JSON with enhanced simulation data")
                    else:
                        logger.debug(f"No additional data extracted from files for {framework}")

                if framework == "pymdp":
                    sim_dir = impl_specific_dir.parent / "simulation_data"
                    sr_candidates = list(sim_dir.glob("*simulation_results.json"))
                    if not sr_candidates and (sim_dir / "simulation_results.json").exists():
                        sr_candidates = [sim_dir / "simulation_results.json"]
                    if sr_candidates:
                        try:
                            with open(sr_candidates[0], encoding="utf-8") as sf:
                                payload = json.load(sf)
                            n_steps = payload.get("num_timesteps")
                            if n_steps is None:
                                n_steps = len(payload.get("observations", []))
                            logger.info(
                                "pymdp_execution_summary model=%s script=%s simulation_results=%s timesteps=%s",
                                model_name,
                                script_info["name"],
                                sr_candidates[0],
                                n_steps,
                            )
                        except (OSError, json.JSONDecodeError, TypeError) as ex:
                            logger.debug("pymdp_execution_summary skipped: %s", ex)

            except Exception as e:
                logger.warning(f"Failed to collect execution outputs: {e}")
                import traceback
                logger.debug(traceback.format_exc())

    except subprocess.TimeoutExpired:
        exec_result['error'] = f'Script execution timed out ({timeout} seconds)'
        logger.error(f"Script {script_info['name']} timed out")

    except Exception as e:
        exec_result['error'] = str(e)
        exec_result['error_type'] = type(e).__name__
        logger.error(f"Error executing {script_info['name']}: {e}")

    return exec_result


# --- Re-export everything from sub-modules for backward compatibility ---
from .data_extractors import (
    collect_execution_outputs,
)
from .data_extractors import (
    extract_simulation_data as _extract_simulation_data,
)
from .data_extractors import (
    extract_simulation_data_from_files as _extract_simulation_data_from_files,
)


def generate_execution_report(execution_results: Dict[str, Any], results_dir: Path, logger: logging.Logger) -> None:
    """
    Generate a comprehensive execution report.
    
    Args:
        execution_results: Dictionary with execution results
        results_dir: Directory to save the report
        logger: Logger instance
    """
    summaries_dir = results_dir / "summaries"
    summaries_dir.mkdir(parents=True, exist_ok=True)
    report_file = summaries_dir / "execution_report.md"

    try:
        with open(report_file, 'w') as f:
            f.write("# GNN Script Execution Report\n\n")
            f.write(f"**Generated:** {execution_results['timestamp']}\n")
            f.write(f"**Target Directory:** {execution_results['target_directory']}\n")
            f.write(f"**Output Directory:** {execution_results['output_directory']}\n\n")

            f.write("## Summary\n\n")
            f.write(f"- **Total Scripts Found:** {execution_results['total_scripts_found']}\n")
            f.write(f"- **Successful Executions:** {execution_results['successful_executions']}\n")
            f.write(f"- **Failed Executions:** {execution_results['failed_executions']}\n")
            skipped = execution_results.get('skipped_executions', 0)
            if skipped:
                f.write(f"- **Skipped (dependency not installed):** {skipped}\n")
            f.write("\n")

            if execution_results['execution_details']:
                f.write("## Execution Details\n\n")

                for detail in execution_results['execution_details']:
                    if detail.get('skipped'):
                        status = "⏭️ SKIPPED"
                    elif detail['success']:
                        status = "✅ SUCCESS"
                    else:
                        status = "❌ FAILED"
                    f.write(f"### {detail['script_name']} - {status}\n\n")
                    f.write(f"- **Framework:** {detail['framework']}\n")
                    f.write(f"- **Executor:** {detail['executor']}\n")
                    f.write(f"- **Path:** `{detail['script_path']}`\n")
                    if detail.get('skipped'):
                        f.write(f"- **Reason:** {detail.get('error', 'Dependency not installed')}\n")
                    else:
                        f.write(f"- **Return Code:** {detail.get('return_code', 'N/A')}\n")
                        f.write(f"- **Execution Time:** {detail.get('execution_time', 0):.2f} seconds\n")

                        if not detail['success'] and 'error' in detail:
                            f.write(f"- **Error:** {detail['error']}\n")

                        if 'output_file' in detail:
                            f.write(f"- **Detailed Output:** {detail['output_file']}\n")

                    f.write("\n")

            f.write("## Next Steps\n\n")
            if execution_results['failed_executions'] > 0:
                f.write("1. Review failed executions above\n")
                f.write("2. Check individual output files for detailed error information\n")
                f.write("3. Ensure required dependencies are installed\n")
                f.write("4. Verify script syntax and functionality\n\n")
            elif skipped:
                f.write("Skipped scripts are due to missing optional dependencies. To run all frameworks: `uv sync --extra execution-frameworks` (or install active-inference, probabilistic-programming, ml-ai, graphs).\n\n")
            else:
                f.write("All scripts executed successfully! Check individual output files for results.\n\n")

        logger.info(f"Generated execution report: {report_file}")

    except Exception as e:
        logger.error(f"Failed to generate execution report: {e}")


def execute_simulation_from_gnn(gnn_file: Path, output_dir: Path) -> Dict[str, Any]:
    """
    Execute simulation from GNN file.
    
    Args:
        gnn_file: Path to GNN file
        output_dir: Output directory
        
    Returns:
        Dictionary with execution results
    """
    try:
        logger.info(f"Executing simulation for {gnn_file}")

        from .executor import GNNExecutor

        engine = GNNExecutor()

        # Execute simulation
        result = engine.execute_simulation_from_gnn(gnn_file, output_dir)

        return result

    except Exception as e:
        logger.error(f"Failed to execute simulation for {gnn_file}: {e}")
        return {
            "success": False,
            "error": str(e)
        }