"""
Main Orchestrator for the Trading Pipeline.

Coordinates data loading, feature engineering, model training, calibration, 
SAC training, and backtesting.
"""

import os
import sys
import logging
import yaml
import pandas as pd
import numpy as np
from datetime import datetime, timezone, timedelta
import argparse
import joblib
import json
from typing import Optional, Any, List, Tuple, Iterator, Dict
import matplotlib.pyplot as plt
import seaborn as sns
import torch # Added for SAC weight aggregation
from collections import OrderedDict # Added for SAC weight aggregation
import time
import shutil
import copy # Added import

# Determine the project root directory based on the script location
# This assumes the script is in src/ and the project root is two levels up
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(script_dir))
# Add project root to sys.path to allow absolute imports from the package
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# Now use absolute imports based on the package structure
from gru_sac_predictor.src.data_loader import DataLoader
# Import other components as they are created
from gru_sac_predictor.src.feature_engineer import FeatureEngineer
# Try importing minimal_whitelist from features.py (assuming it exists there)
try:
    from gru_sac_predictor.src.features import minimal_whitelist
except ImportError:
    # Fallback: Define it here if features.py doesn't exist or doesn't have it
    logging.warning("Could not import minimal_whitelist from .features, defining fallback.")
    minimal_whitelist = [
        "return_1m", "return_15m", "return_60m", "ATR_14", "volatility_14d",
        "chaikin_AD_10", "svi_10", "EMA_10", "EMA_50", "MACD", "MACD_signal",
        "hour_sin", "hour_cos",
    ]
from gru_sac_predictor.src.gru_model_handler import GRUModelHandler
from gru_sac_predictor.src.calibrator import Calibrator
# --- Import Vector Calibrator (Task 4) --- #
try:
    from gru_sac_predictor.src.calibrator_vector import VectorCalibrator, VECTOR_CALIBRATOR_AVAILABLE
except ImportError:
    logging.warning("VectorCalibrator could not be imported. Vector scaling method will not be available.")
    VectorCalibrator = None # Define as None if import fails
    VECTOR_CALIBRATOR_AVAILABLE = False
# --- End Import --- #
from gru_sac_predictor.src.sac_trainer import SACTrainer
from gru_sac_predictor.src.sac_agent import SACTradingAgent # Added for SAC agent loading
from gru_sac_predictor.src.backtester import Backtester
from gru_sac_predictor.src.baseline_checker import BaselineChecker # Import BaselineChecker
from gru_sac_predictor.src.metrics import edge_filtered_accuracy, calculate_brier_score, calculate_sharpe_ratio, _calculate_optimal_edge_threshold # Added calculate_sharpe_ratio if needed for aggregation

# Removed redundant imports for feature selection
from sklearn.preprocessing import StandardScaler
# --- Add imports for baseline --- #
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, roc_auc_score, classification_report
from sklearn.model_selection import train_test_split
import scipy.stats as st
# --- Import edge_filtered_accuracy (Task 6.1/6.2) --- #
try:
     # Ensure both metrics are imported
     from gru_sac_predictor.src.metrics import edge_filtered_accuracy, calculate_brier_score, _calculate_optimal_edge_threshold
except ImportError:
     logging.error("Failed to import metrics from gru_sac_predictor.src.metrics. Validation check will fail.")
     # Define placeholders
     def edge_filtered_accuracy(*args, **kwargs): return np.nan, 0
     def calculate_brier_score(*args, **kwargs): return np.nan
     def _calculate_optimal_edge_threshold(*args, **kwargs): return 0.1 # Default fallback
# --- End Import --- #
# --- End imports for baseline --- #

# --- Import Stage Functions --- #
from gru_sac_predictor.src.pipeline_stages.data_processing import (
    engineer_features_for_fold,
    define_labels_and_align_fold,
    split_data_fold
)
from gru_sac_predictor.src.pipeline_stages.feature_processing import (
    scale_features_fold,
    select_features_fold,
    prune_features_fold # Added import
)
from gru_sac_predictor.src.pipeline_stages.sequence_creation import (
    create_sequences_fold # Added import
)
from gru_sac_predictor.src.pipeline_stages.evaluation import run_baseline_checks_fold # Added baseline check import
from gru_sac_predictor.src.pipeline_stages.evaluation import run_gru_validation_checks_fold # Added validation check import
from gru_sac_predictor.src.pipeline_stages.evaluation import run_backtest_fold # Import the new backtest function
from gru_sac_predictor.src.pipeline_stages.modelling import train_or_load_gru_fold
from gru_sac_predictor.src.pipeline_stages.modelling import calibrate_probabilities_fold
from gru_sac_predictor.src.pipeline_stages.modelling import train_or_load_sac_fold
from gru_sac_predictor.src.pipeline_stages.modelling import aggregate_sac_agents

# --- Import FoldGenerator --- #
from gru_sac_predictor.src.fold_generator import FoldGenerator

logger = logging.getLogger(__name__) # Use module-level logger

# --- Refactored Label Generation Logic --- #
# [Function _generate_direction_labels removed - Moved to data_processing.py]
# --- End Refactored Label Generation --- #

# --- Gap & Regime Helper Functions --- #
def add_regime_tags(df: pd.DataFrame, indicator: str, window: int, quantiles: List[float]) -> pd.DataFrame:
    """
    Calculates a regime indicator and adds a 'regime_tag' column.

    Args:
        df (pd.DataFrame): Input dataframe with features (must include 'close' or 'return').
        indicator (str): The type of indicator ('volatility').
        window (int): The rolling window size for the indicator.
        quantiles (List[float]): Quantiles to define regime boundaries.

    Returns:
        pd.DataFrame: DataFrame with 'regime_tag' column added.
    """
    df = df.copy()
    if indicator == 'volatility':
        # Calculate rolling volatility (std dev of log returns)
        if 'return_1m' not in df.columns: # Assuming 'return_1m' or similar exists
            df['log_ret'] = np.log(df['close'] / df['close'].shift(1))
            indicator_series = df['log_ret'].rolling(window=window, min_periods=window // 2).std()
            del df['log_ret']
        else:
            # Use a pre-calculated short-term return if available
            indicator_series = df['return_1m'].rolling(window=window, min_periods=window // 2).std()
    # Add other indicator calculations here (e.g., trend, RSI)
    else:
        logger.warning(f"Unsupported regime indicator: {indicator}. Skipping regime tagging.")
        df['regime_tag'] = -1 # Assign a default tag
        return df

    # Calculate quantile thresholds
    thresholds = indicator_series.quantile(quantiles).tolist()
    
    # Assign regime tags based on thresholds
    def assign_tag(value):
        if pd.isna(value):
            return -1 # Handle NaNs (e.g., at the start of the series)
        for i, threshold in enumerate(thresholds):
            if value <= threshold:
                return i
        return len(thresholds) # Tag for values above the highest threshold

    df['regime_tag'] = indicator_series.apply(assign_tag)
    logger.info(f"Regime tags added using '{indicator}' (window={window}). Distribution:\n{df['regime_tag'].value_counts(normalize=True).sort_index()}")
    return df

def split_into_contiguous_chunks(df: pd.DataFrame, gap_threshold_minutes: int) -> List[Tuple[pd.Timestamp, pd.Timestamp]]:
    """
    Splits a DataFrame index into contiguous chunks based on time gaps.

    Args:
        df (pd.DataFrame): DataFrame with a DatetimeIndex.
        gap_threshold_minutes (int): Maximum allowed gap in minutes before splitting.

    Returns:
        List[Tuple[pd.Timestamp, pd.Timestamp]]: List of (start, end) timestamps for each chunk.
    """
    if not isinstance(df.index, pd.DatetimeIndex):
        raise ValueError("DataFrame must have a DatetimeIndex.")
    
    if df.empty:
        return []

    chunks = []
    start_time = df.index[0]
    gap_threshold = pd.Timedelta(minutes=gap_threshold_minutes)

    time_diffs = df.index.to_series().diff()
    split_points = time_diffs[time_diffs > gap_threshold].index

    last_end_time = start_time
    for split_time in split_points:
        # Find the index entry *before* the split point
        split_idx_loc = df.index.get_loc(split_time)
        if split_idx_loc > 0:
            end_time = df.index[split_idx_loc - 1]
            if end_time >= start_time:
                chunks.append((start_time, end_time))
                logger.debug(f"Detected chunk: [{start_time}, {end_time}] due to gap before {split_time}")
            # Start the next chunk at the split_time itself
            start_time = split_time 
            last_end_time = end_time # Keep track for the final chunk
    
    # Add the final chunk
    final_end_time = df.index[-1]
    if final_end_time >= start_time:
         chunks.append((start_time, final_end_time))
         logger.debug(f"Adding final chunk: [{start_time}, {final_end_time}]")

    logger.info(f"Split data into {len(chunks)} contiguous chunks based on gap threshold > {gap_threshold_minutes} minutes.")
    return chunks
# --- End Gap & Regime Helpers --- #

class TradingPipeline:
    """Orchestrates the entire trading strategy pipeline."""

    def __init__(self, config: dict, io_manager: Optional[Any] = None):
        """
        Initialize the pipeline with configuration, optional CLI args, and IOManager.
        
        Args:
            config (dict): The loaded configuration dictionary.
            io_manager (IOManager, optional): Initialized IOManager instance. Defaults to None.
        """
        # Store the passed config dictionary directly
        self.config = config
        # Run ID and Git SHA should be generated *before* logger/io setup in run.py
        # and passed via the IOManager.
        if io_manager is None:
            # IOManager is considered essential for proper operation.
            # Raise an error or handle appropriately if not provided.
            # For now, log critical error and exit, assuming IOManager is required.
            # TODO: Decide final handling if IOManager *can* be optional.
            logging.critical("IOManager not provided during TradingPipeline initialization. Cannot proceed.")
            raise ValueError("IOManager instance is required for TradingPipeline.")
            # Fallback removed - rely on IOManager
            # self.run_id = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S_fallback")
            # self.git_sha = "unknown"
            # logger_to_use = logging # Use root logger if no io/logger setup provided
        else:
             # --- Retrieve run_id and git_sha FROM io_manager ---
             if not hasattr(io_manager, 'run_id') or not io_manager.run_id:
                 raise ValueError("IOManager instance provided, but does not have a 'run_id' attribute set.")
             self.run_id = io_manager.run_id

             # Assume git_sha is also an attribute of the initialized IOManager
             if not hasattr(io_manager, 'git_sha') or not io_manager.git_sha:
                 # logging.warning("IOManager instance does not have 'git_sha' attribute. Using 'unknown'.")
                 # self.git_sha = "unknown" # Or raise error if mandatory
                 # For now, let's assume it's mandatory for traceability
                 raise ValueError("IOManager instance provided, but does not have a 'git_sha' attribute set.")
             self.git_sha = io_manager.git_sha
             # --- End Retrieval ---
             
             logger_to_use = logging.getLogger() # Assume logger was set up by IOManager/run.py
             
        # self.io is now guaranteed to be an IOManager instance if we proceed past checks
        self.io = io_manager 
        # Store git_sha on self as well, retrieved via IOManager
        # Ensure git_sha exists on io_manager (already checked above, but good practice)
        self.git_sha = getattr(io_manager, 'git_sha', 'unknown_in_pipeline') 
        self.pipeline_version = "3.0.0" # Placeholder version

        # --- Directory Setup (Now handled by IOManager if provided) --- #
        if self.io:
            self.dirs = {
                 'results': self.io.run_results_dir,
                 'models': self.io.run_models_dir,
                 'logs': self.io.run_logs_dir,
                 'figures': self.io.run_figures_dir # Added figures dir
            }
            self.base_models_dir_path = self.io.models_dir # Base models dir (parent of run dir)
            self.current_run_models_dir = self.io.run_models_dir
            # Logging setup is now done *before* pipeline init in run.py
            # self._setup_logging() # Remove internal logging setup
        else:
             logger_to_use.warning("IOManager not provided. Setting up directories manually.")
             # Fallback manual setup (might be removed if IOManager is required)
             self._setup_directories_manual() 
             self._setup_logging_manual() # Fallback logging
        # --- End Directory Setup --- #

        # --- Feature Whitelist Modification --- #
        # Add 'bar_imputed' to the minimal whitelist if it's not already there
        if 'bar_imputed' not in minimal_whitelist:
             minimal_whitelist.append('bar_imputed')
             logger.info("Added 'bar_imputed' to minimal_whitelist.")
        # --- End Whitelist Modification --- #

        # Log Banner (Moved to run.py which has version info)
        # logger_to_use.info(...) 

        # --- Initialize Components --- #
        # Extract db_directory from config before passing to DataLoader
        data_cfg = self.config.get('data', {})
        db_directory_path = data_cfg.get('db_dir', 'data/db') # Use 'db_dir' key, provide a default
        if not db_directory_path or not isinstance(db_directory_path, str):
             # Adjust error message to reflect the correct key 'db_dir'
             logger.error(f"Invalid or missing 'db_dir' in config['data']. Found: {db_directory_path}. Using default 'data/db'.")
             db_directory_path = 'data/db' # Fallback

        self.data_loader = DataLoader(db_dir=db_directory_path)
        self.feature_engineer = FeatureEngineer(self.config)
        # Extract edge threshold for Calibrator initialization
        calibration_cfg = self.config.get('calibration', {})
        initial_edge_threshold = calibration_cfg.get('edge_threshold', 0.1) # Get edge from config
        self.calibrator = Calibrator(edge_threshold=initial_edge_threshold)
        # --- Vector Calibrator (Task 4) --- #
        if VECTOR_CALIBRATOR_AVAILABLE:
             self.vector_calibrator = VectorCalibrator() # Initialize without config
        else:
             self.vector_calibrator = None
        # --- End Vector Calibrator --- #
        # Initialize SACTrainer only when needed (in train_or_load_sac)
        # self.sac_trainer = SACTrainer(config=self.config) 
        self.backtester = Backtester(self.config, io_manager=self.io) # Pass io_manager
        # Initialize gru_handler (needs run_id and models dir)
        self.gru_handler = GRUModelHandler(
            run_id=self.run_id, 
            models_dir=self.current_run_models_dir,
            config=self.config 
        )
        # Initialize BaselineChecker
        self.baseline_checker = BaselineChecker(self.config, io=self.io) # Pass io
        # --- End Initialize Components --- #

        # --- Initialize state variables --- #
        self.df_raw = None
        self.load_summary = None # Store load summary
        self.df_engineered_full = None
        # self.df_features_minimal = None # Removed minimal pruning here
        self.df_labeled_aligned = None
        self.X_raw_aligned = None
        self.y_aligned = None   
        self.y_dir_aligned = None
        self.X_train_raw = None
        self.X_val_raw = None
        self.X_test_raw = None
        self.y_train = None
        self.y_val = None
        self.y_test = None
        self.df_train_original = None # Store original data for splits
        self.df_val_original = None
        self.df_test_original = None
        self.y_dir_train = None
        self.scaler = None
        self.X_train_scaled = None
        self.X_val_scaled = None
        self.X_test_scaled = None
        self.final_whitelist = None
        self.X_train_pruned = None
        self.X_val_pruned = None
        self.X_test_pruned = None
        self.X_train_seq = None
        self.X_val_seq = None
        self.X_test_seq = None
        self.y_train_seq_dict = None
        self.y_val_seq_dict = None
        self.y_test_seq_dict = None
        self.train_indices = None
        self.val_indices = None
        self.test_indices = None
        self.gru_model = None
        self.gru_model_run_id_loaded_from = None # Track which run ID model came from
        self.optimal_T = None
        self.vector_cal_params = None # Store vector calibration parameters
        self.sac_agent_load_path = None # Path to the SAC agent to load for backtesting
        self.backtest_results_df = None
        self.backtest_metrics = None
        self.metrics_log_df = None # For logging detailed metrics
        self.use_ternary = self.config.get('gru', {}).get('use_ternary', False) # Cache ternary flag
        self.aggregated_metrics: Optional[dict] = None # Aggregated metrics across folds
        self.optimized_edge_threshold: Optional[float] = None # Store optimized edge threshold per fold
        # --- Add attributes for baseline filtering --- #
        self.fwd_returns_aligned: Optional[pd.Series] = None
        self.eps_aligned: Optional[pd.Series] = None
        # Attributes for split returns/eps needed by baseline check
        self.fwd_ret_train: Optional[pd.Series] = None
        self.eps_train: Optional[pd.Series] = None
        self.fwd_ret_val: Optional[pd.Series] = None
        self.eps_val: Optional[pd.Series] = None
        self.y_dir_val_ordinal: Optional[pd.Series] = None # <<< ADDED
        # --- Feature selection optimization --- #
        self.initial_l1_whitelist: Optional[List[str]] = None
        # --- End Add --- #
        # --- End Initialize state variables --- #
        
        # Save config handled by run.py via IOManager typically
        # self._save_run_config() 
        if self.io:
             # Maybe save it again here for completeness? Or rely on run.py?
             config_save_path = self.io.path('results', 'run_config', suffix='.yaml')
             try:
                  with open(config_save_path, 'w') as f:
                       yaml.dump(self.config, f, default_flow_style=False)
                  logger_to_use.info(f"Saved run configuration copy to {config_save_path}")
             except Exception as e:
                  logger_to_use.error(f"Failed to save run configuration via IOManager: {e}")
        else:
             logger_to_use.warning("IOManager not available, cannot save run config copy from pipeline.")

        self.fold_generator = None # Initialize fold generator attribute

    # --- Remove Wrapper Method Definitions --- #
    # def define_labels_and_align(self, df_engineered: pd.DataFrame) -> Tuple[pd.DataFrame, str, List[str]]:
    #     # ... implementation removed ...
    #     pass

    # def split_data(self, df_labeled_aligned_fold: pd.DataFrame, fold_dates: Optional[Tuple] = None):
    #     # ... implementation removed ...
    #     pass

    # def run_baseline_checks(self, fold_num: int, X_train_to_use: Optional[pd.DataFrame] = None, X_val_to_use: Optional[pd.DataFrame] = None) -> Dict[str, Any]:
    #     # ... implementation removed ...
    #     pass
    # --- End Remove Wrapper Methods --- #

    # --- execute method refactored for Walk-Forward --- #
    def execute(self):
        """Runs the full trading pipeline, potentially using Walk-Forward validation."""
        logger.info(f"--- Starting Trading Pipeline: Run ID {self.run_id} ---")

        # 1. Load and Preprocess FULL Data
        # Replace direct call with call to data_loader using config
        data_cfg = self.config.get('data', {})
        ticker = data_cfg.get('ticker')
        exchange = data_cfg.get('exchange')
        start_date = data_cfg.get('start_date')
        end_date = data_cfg.get('end_date')
        interval = data_cfg.get('interval', '1min') # Default to 1min if not specified
        vol_sampling_cfg = data_cfg.get('volatility_sampling', {})
        vol_sampling_enabled = vol_sampling_cfg.get('enabled', False)
        vol_window = vol_sampling_cfg.get('window', 30)
        vol_quantile = vol_sampling_cfg.get('quantile', 0.5)

        if not all([ticker, exchange, start_date, end_date]):
             logger.critical("Missing required data parameters (ticker, exchange, start_date, end_date) in config['data']. Cannot load data.")
             raise ValueError("Missing required data configuration.")
             
        try:
            logger.info(f"Loading data for {ticker} from {exchange} between {start_date} and {end_date} at {interval} interval.")
            # Call the load_data method on the data_loader instance
            self.df_raw = self.data_loader.load_data(
                ticker=ticker,
                exchange=exchange,
                start_date=start_date,
                end_date=end_date,
                interval=interval,
                vol_sampling=vol_sampling_enabled,
                vol_window=vol_window,
                vol_quantile=vol_quantile
            )
            # The load_data method in DataLoader now includes imputation and reporting

        except Exception as e:
            logger.critical(f"Failed to load initial data: {e}", exc_info=True)
            self.df_raw = None
            raise SystemExit("Failed to load initial data.") from e

        if self.df_raw is None or self.df_raw.empty:
            logger.error("Initial data loading and preprocessing resulted in an empty DataFrame. Exiting pipeline.")
            if hasattr(self.data_loader, 'last_load_summary') and self.data_loader.last_load_summary:
                logger.error(f"Load Summary: {self.data_loader.last_load_summary}")
            raise SystemExit("Data loading failed.")

        # --- Initialize FoldGenerator AFTER data is loaded --- #
        try:
             logger.info("Initializing FoldGenerator...")
             self.fold_generator = FoldGenerator(config=self.config, df_raw=self.df_raw, io_manager=self.io)
             logger.info("FoldGenerator initialized.")
        except Exception as fg_init_err:
             logger.error(f"Failed to initialize FoldGenerator: {fg_init_err}", exc_info=True)
             raise SystemExit("FoldGenerator initialization failed.") from fg_init_err
        # --- END FoldGenerator Initialization --- #

        # 2. Generate Walk-Forward Folds using the generator instance
        try:
             # Make generate_folds return an iterator/generator
             fold_dates_iterator = self.fold_generator.generate_folds()
        except Exception as fg_gen_err:
             logger.error(f"Error calling fold_generator.generate_folds(): {fg_gen_err}", exc_info=True)
             raise SystemExit("Fold generation failed.") from fg_gen_err

        self.all_fold_metrics = [] # Reset fold metrics list for the run
        all_successful_sac_agent_paths = [] # Store paths of successfully trained SAC agents per fold
        fold_count = 0

        # 3. Loop Through Folds
        for fold_dates in fold_dates_iterator:
            fold_count += 1
            self.current_fold = fold_count # Set current fold number
            logger.info(f"=== Processing Fold {self.current_fold} ===")
            self.current_fold_report_data = [] # <<< RE-INITIALIZE report list for the fold >>>

            # --- Handle single split case is now handled INTERNALLY by FoldGenerator ---
            # No need for the explicit single split date calculation here anymore.
            # The generator yields the correct dates directly.
            
            if fold_dates is None or len(fold_dates) != 6:
                 logger.error(f"Fold {self.current_fold}: Invalid fold dates received from generator: {fold_dates}. Skipping.")
                 self.all_fold_metrics.append({'fold': self.current_fold, 'status': 'error', 'reason': 'Invalid fold dates received'})
                 continue

            train_start, train_end, val_start, val_end, test_start, test_end = fold_dates
            
            # Setup fold-specific directories using IOManager if available
            if self.io:
                 # self.fold_dirs = self.io.setup_fold_dirs(self.current_fold)
                 self.fold_dirs = self.io.get_fold_dirs(self.current_fold) # Correct method name
            else:
                 self.fold_dirs = {} # Set empty if no IOManager
                 logger.warning(f"Fold {self.current_fold}: IOManager not available, cannot create fold-specific directories.")

            # Select data for the current fold using the ORIGINAL df_raw
            # The generator used df_raw_tagged internally, but pipeline stages use df_raw
            fold_start_date = train_start
            # Determine the latest end date actually provided for slicing
            fold_end_candidates = [d for d in [test_end, val_end, train_end] if pd.notna(d)]
            if not fold_end_candidates:
                 logger.error(f"Fold {self.current_fold}: No valid end date found in fold dates: {fold_dates}. Skipping fold.")
                 self.all_fold_metrics.append({'fold': self.current_fold, 'status': 'error', 'reason': 'No valid end date'})
                 continue
            fold_end_date = max(fold_end_candidates)

            # Ensure start/end dates have timezone consistent with df_raw index if applicable
            if self.df_raw.index.tz is not None:
                 if pd.notna(fold_start_date) and fold_start_date.tzinfo is None:
                     fold_start_date = fold_start_date.tz_localize(self.df_raw.index.tz)
                 if pd.notna(fold_end_date) and fold_end_date.tzinfo is None:
                     fold_end_date = fold_end_date.tz_localize(self.df_raw.index.tz)
            
            # Check validity before slicing
            if pd.isna(fold_start_date) or pd.isna(fold_end_date) or fold_end_date < fold_start_date:
                 logger.error(f"Fold {self.current_fold}: Invalid start/end range for slicing df_raw: Start={fold_start_date}, End={fold_end_date}. Skipping fold.")
                 self.all_fold_metrics.append({'fold': self.current_fold, 'status': 'error', 'reason': 'Invalid date range for fold data selection'})
                 continue
                 
            try:
                current_fold_data_raw = self.df_raw.loc[fold_start_date:fold_end_date]
            except Exception as slice_err:
                 logger.error(f"Fold {self.current_fold}: Error slicing df_raw for dates [{fold_start_date} to {fold_end_date}]: {slice_err}. Skipping fold.", exc_info=True)
                 self.all_fold_metrics.append({'fold': self.current_fold, 'status': 'error', 'reason': f'Error slicing raw data: {slice_err}'})
                 continue

            if current_fold_data_raw.empty:
                 logger.warning(f"Fold {self.current_fold}: No raw data found for range {fold_start_date} to {fold_end_date}. Skipping fold.")
                 # Store failure info
                 self.all_fold_metrics.append({'fold': self.current_fold, 'status': 'skipped', 'reason': 'No raw data in date range'})
                 continue
            
            logger.info(f"Fold {self.current_fold}: Raw data range for pipeline stages [{current_fold_data_raw.index.min()}, {current_fold_data_raw.index.max()}]")

            # --- Run Pipeline Steps within the Fold --- #
            try:
                # a. Engineer Features (using stage function)
                df_engineered_fold = engineer_features_for_fold(
                    df=current_fold_data_raw, 
                    feature_engineer=self.feature_engineer,
                    io=self.io, 
                    config=self.config,
                    target_col=None 
                )
                if df_engineered_fold.empty:
                     raise SystemExit(f"Fold {self.current_fold}: Feature engineering resulted in empty dataframe.")

                # b. Define Labels and Align (Direct stage function call)
                logger.info(f"Fold {self.current_fold}: Calling Stage: Defining Labels and Aligning...")
                df_labeled_aligned_fold, target_dir_col, target_cols, fwd_returns_aligned, eps_aligned = define_labels_and_align_fold(
                    df_engineered=df_engineered_fold,
                    config=self.config
                )
                if df_labeled_aligned_fold.empty:
                    raise SystemExit(f"Fold {self.current_fold}: Label definition stage resulted in empty dataframe.")
                # Store results on self
                self.target_dir_col = target_dir_col
                self.target_columns = target_cols
                self.fwd_returns_aligned = fwd_returns_aligned
                self.eps_aligned = eps_aligned

                # c. Split data (Direct stage function call)
                logger.info(f"Fold {self.current_fold}: Calling Stage: Splitting Data...")
                (
                    self.X_train_raw, self.X_val_raw, self.X_test_raw,
                    self.y_train, self.y_val, self.y_test,
                    self.df_train_original, self.df_val_original, self.df_test_original, 
                    self.y_dir_train_ordinal,
                    self.fwd_ret_train, self.fwd_ret_val,
                    self.eps_train, self.eps_val,
                    self.y_dir_val_ordinal # Captures val ordinal dir labels
                ) = split_data_fold(
                    df_labeled_aligned=df_labeled_aligned_fold,
                    fwd_returns_aligned=self.fwd_returns_aligned,
                    eps_aligned=self.eps_aligned,
                    config=self.config,
                    target_columns=self.target_columns, 
                    target_dir_col=self.target_dir_col, 
                    fold_dates=fold_dates,
                    current_fold=self.current_fold,
                    io=self.io 
                )
                # Note: split_data_fold raises SystemExit on failure
                logger.info(f"Fold {self.current_fold}: Data splitting stage complete.")

                # d. Scale Features (using stage function)
                (self.X_train_scaled, 
                 self.X_val_scaled, 
                 self.X_test_scaled, 
                 self.scaler) = scale_features_fold(
                    X_train_raw=self.X_train_raw,
                    X_val_raw=self.X_val_raw,
                    X_test_raw=self.X_test_raw,
                    run_id=self.run_id,
                    fold_num=self.current_fold,
                    fold_models_dir=self.fold_dirs.get('models', '.'),
                    main_run_models_dir=self.current_run_models_dir,
                )
                # Note: scale_features_fold now raises SystemExit on failure

                # <<< NEW: Coarse Univariate Filter >>>
                X_train_coarse, X_val_coarse = None, None
                feat_cfg = self.config.get('features', {})
                coarse_quantile = feat_cfg.get('coarse_univariate_quantile', 0.70)
                logger.info(f"Fold {self.current_fold}: Applying coarse univariate filter (keeping top {1-coarse_quantile:.0%}) before first baseline check...")
                try:
                    if self.y_dir_train_ordinal is not None and self.X_train_scaled is not None:
                        y_aligned, X_aligned = self.y_dir_train_ordinal.align(self.X_train_scaled, join='inner', axis=0)
                        y_numeric = pd.to_numeric(y_aligned, errors='coerce').fillna(0)
                        X_train_aligned_for_coarse = X_aligned
                        
                        if X_train_aligned_for_coarse.empty:
                            logger.warning(f"Fold {self.current_fold}: X_train became empty after alignment for coarse filter. Skipping filter.")
                        else:
                            corrs = X_train_aligned_for_coarse.apply(lambda col: abs(np.corrcoef(col, y_numeric)[0, 1]) if col.var() > 1e-6 else 0.0)
                            corrs.fillna(0, inplace=True)
                            if not corrs.empty and not corrs.isnull().all():
                                quantile_value = corrs.quantile(coarse_quantile)
                                keep_mask = corrs >= quantile_value
                                keep_coarse_features = corrs[keep_mask].index.tolist()
                                if keep_coarse_features:
                                    X_train_coarse = self.X_train_scaled[keep_coarse_features]
                                    if self.X_val_scaled is not None:
                                        X_val_coarse = self.X_val_scaled[[col for col in keep_coarse_features if col in self.X_val_scaled.columns]]
                                    logger.info(f"Fold {self.current_fold}: Coarse filter kept {len(keep_coarse_features)} features (Corr >= {quantile_value:.4f}).")
                                else:
                                    logger.warning(f"Fold {self.current_fold}: Coarse univariate filter removed all features. Baseline will run on full scaled set.")
                            else:
                                 logger.warning(f"Fold {self.current_fold}: Could not compute valid correlations for coarse filter. Baseline will run on full scaled set.")
                    else:
                        logger.warning(f"Fold {self.current_fold}: Skipping coarse filter due to missing y_dir_train_ordinal or X_train_scaled.")
                except Exception as coarse_err:
                     logger.error(f"Fold {self.current_fold}: Error during coarse univariate filtering: {coarse_err}. Baseline will run on full scaled set.", exc_info=True)
                # <<< END NEW Filter >>>

                # e. Baseline Check (Direct stage function call)
                # --- Select Input Data for Baseline --- #
                run_on_coarse = X_train_coarse is not None
                X_train_for_check = X_train_coarse if run_on_coarse else self.X_train_scaled
                X_val_for_check = X_val_coarse if run_on_coarse else self.X_val_scaled
                log_msg_data_source = "COARSELY FILTERED" if run_on_coarse else "FULL SCALED"
                logger.info(f"Fold {self.current_fold}: Running baseline checks on {log_msg_data_source} data.")
                # --- End Input Data Selection --- #
                
                # Check prerequisites before calling stage function
                baseline_prereqs_met = True
                required_attrs_baseline = [
                     'y_dir_train_ordinal', 'y_dir_val_ordinal',
                     'fwd_ret_train', 'fwd_ret_val',
                     'eps_train', 'eps_val',
                     'baseline_checker'
                ]
                missing_attrs = [attr for attr in required_attrs_baseline if not hasattr(self, attr) or getattr(self, attr) is None]
                if X_train_for_check is None or X_val_for_check is None:
                    missing_attrs.append("scaled_features (or coarse features)")
                
                if missing_attrs:
                     logger.error(f"Fold {self.current_fold}: Cannot run baseline checks. Missing required attributes/data: {missing_attrs}")
                     baseline_prereqs_met = False
                     baseline_report = { # Manually create error report
                        "gate_name": "baseline_checks",
                        "status": "error",
                        "timestamp": datetime.now(timezone.utc).isoformat(),
                        "reason": f"Missing required inputs: {missing_attrs}",
                        "details": {}
                     }
                
                if baseline_prereqs_met:
                    baseline_report = run_baseline_checks_fold(
                        X_train_scaled=X_train_for_check,
                        X_val_scaled=X_val_for_check,
                        y_train_dir_ordinal=self.y_dir_train_ordinal, 
                        y_val_dir_ordinal=self.y_dir_val_ordinal, 
                        fwd_ret_train=self.fwd_ret_train,
                        eps_train=self.eps_train,
                        fwd_ret_val=self.fwd_ret_val,
                        eps_val=self.eps_val,
                        baseline_checker=self.baseline_checker,
                        config=self.config,
                        io=self.io,
                        fold_num=self.current_fold,
                        fold_dirs=self.fold_dirs,
                        base_results_dir=self.io.run_results_dir if self.io else '.' 
                    )
                
                # --- Report Handling (same as before) --- #
                self.current_fold_report_data.append(baseline_report)
                if baseline_report.get('status') in ['failed', 'error']:
                     raise SystemExit(f"Fold {self.current_fold} failed baseline checks: {baseline_report.get('reason')}")
                # --- End report handling --- #

                # f. Select Features & Prune (using stage functions)
                logger.info(f"Fold {self.current_fold}: Selecting features...")
                self.final_whitelist = select_features_fold(
                    X_train_raw=self.X_train_raw,
                    y_dir_train_ordinal=self.y_dir_train_ordinal, # Pass ordinal labels
                    feature_engineer=self.feature_engineer,
                    io=self.io,
                    run_id=self.run_id,
                    fold_num=self.current_fold,
                    fold_models_dir=self.fold_dirs.get('models', '.'),
                    fold_results_dir=self.fold_dirs.get('results', '.'),
                    main_run_models_dir=self.current_run_models_dir
                )
                # Note: select_features_fold raises SystemExit on failure
                
                logger.info(f"Fold {self.current_fold}: Pruning scaled features using whitelist ({len(self.final_whitelist)} features)...")
                (self.X_train_pruned, 
                 self.X_val_pruned, 
                 self.X_test_pruned) = prune_features_fold(
                    X_train_scaled=self.X_train_scaled,
                    X_val_scaled=self.X_val_scaled,
                    X_test_scaled=self.X_test_scaled,
                    final_whitelist=self.final_whitelist,
                    feature_engineer=self.feature_engineer, # Pass the instance
                    fold_num=self.current_fold
                )
                # Note: prune_features_fold raises SystemExit on failure

                # <<< Pruning handled within select_and_prune_features >>>
                # <<< The above calls replace the old self.select_and_prune_features() >>>

                # g. Post-Pruning Baseline Check (Direct stage function call)
                # --- Capture report and check status --- #
                # Replace wrapper call with direct stage function call
                post_pruning_baseline_report = run_baseline_checks_fold(
                    X_train_scaled=self.X_train_pruned, # Pass PRUNED data
                    X_val_scaled=self.X_val_pruned,     # Pass PRUNED data
                    y_train_dir_ordinal=self.y_dir_train_ordinal, 
                    y_val_dir_ordinal=self.y_dir_val_ordinal, 
                    fwd_ret_train=self.fwd_ret_train,
                    eps_train=self.eps_train,
                    fwd_ret_val=self.fwd_ret_val,
                    eps_val=self.eps_val,
                    baseline_checker=self.baseline_checker, 
                    config=self.config,
                    io=self.io, 
                    fold_num=self.current_fold,
                    fold_dirs=self.fold_dirs, 
                    base_results_dir=self.io.run_results_dir if self.io else '.' 
                )
                # Rename the gate for clarity in reports
                post_pruning_baseline_report['gate_name'] = 'post_pruning_baseline_checks' 
                self.current_fold_report_data.append(post_pruning_baseline_report)
                if post_pruning_baseline_report.get('status') in ['failed', 'error']:
                     raise SystemExit(f"Fold {self.current_fold} failed post-pruning baseline checks: {post_pruning_baseline_report.get('reason')}")
                # --- End report handling --- #
                
                # h. Update Scaled Data to Pruned Data
                logger.info(f"Fold {self.current_fold}: Post-pruning checks passed. Updating scaled data to pruned versions for subsequent steps.")
                self.X_train_scaled = self.X_train_pruned
                self.X_val_scaled = self.X_val_pruned
                self.X_test_scaled = self.X_test_pruned
                # Clean up pruned attributes to avoid confusion if needed? Optional.
                # del self.X_train_pruned, self.X_val_pruned, self.X_test_pruned
                # <<< End Update >>>

                # i. Create Sequences (using stage function)
                # Replace single method call with calls to stage function for each split
                logger.info(f"Fold {self.current_fold}: Creating sequences...")
                gru_cfg = self.config.get('gru', {})
                lookback = gru_cfg.get('lookback', 60)
                
                if not hasattr(self, 'target_columns') or not self.target_columns:
                     raise SystemExit(f"Fold {self.current_fold}: Target column names not set before create_sequences stage.")
                     
                # --- Create Train Sequences --- #
                (self.X_train_seq, 
                 self.y_train_seq_dict, 
                 self.train_indices, 
                 dropped_train) = create_sequences_fold(
                    X_data=self.X_train_pruned,
                    y_data=self.df_train_original, # Original df contains targets
                    target_names=self.target_columns,
                    lookback=lookback,
                    name="Train",
                    config=self.config,
                    io=self.io
                 )
                if self.X_train_seq is None:
                    raise SystemExit(f"Fold {self.current_fold}: Failed to create training sequences.")
                
                # --- Create Validation Sequences --- #
                (self.X_val_seq, 
                 self.y_val_seq_dict, 
                 self.val_indices, 
                 dropped_val) = create_sequences_fold(
                    X_data=self.X_val_pruned,
                    y_data=self.df_val_original, # Original df contains targets
                    target_names=self.target_columns,
                    lookback=lookback,
                    name="Validation",
                    config=self.config,
                    io=self.io
                 )
                if self.X_val_seq is None:
                    raise SystemExit(f"Fold {self.current_fold}: Failed to create validation sequences.")
                    
                # --- Create Test Sequences --- #
                (self.X_test_seq, 
                 self.y_test_seq_dict, 
                 self.test_indices, 
                 dropped_test) = create_sequences_fold(
                    X_data=self.X_test_pruned,
                    y_data=self.df_test_original, # Original df contains targets
                    target_names=self.target_columns,
                    lookback=lookback,
                    name="Test",
                    config=self.config,
                    io=self.io
                 )
                if self.X_test_seq is None:
                     raise SystemExit(f"Fold {self.current_fold}: Failed to create test sequences.")
                # --- End Capture --- #
                logger.info(f"Fold {self.current_fold}: Sequence creation complete. Dropped counts: Train={dropped_train}, Val={dropped_val}, Test={dropped_test}")

                # j. Train/Load GRU Model (using stage function)
                # Replace method call with stage function call
                logger.info(f"Fold {self.current_fold}: Training or loading GRU model...")
                (
                    self.gru_model,
                    self.gru_handler, # Handler might be updated internally
                    self.gru_model_run_id_loaded_from,
                    scaler_returned, # Scaler associated with the loaded/trained model
                    X_train_seq_new,
                    y_train_seq_dict_new,
                    train_indices_new,
                    X_val_seq_new,
                    y_val_seq_dict_new,
                    val_indices_new,
                    X_test_seq_new,
                    y_test_seq_dict_new,
                    test_indices_new
                 ) = train_or_load_gru_fold(
                    config=self.config,
                    run_id=self.run_id,
                    current_fold=self.current_fold,
                    current_run_models_dir=self.current_run_models_dir,
                    base_models_dir_path=self.base_models_dir_path,
                    gru_handler=self.gru_handler,
                    X_train_seq=self.X_train_seq,
                    y_train_seq_dict=self.y_train_seq_dict,
                    X_val_seq=self.X_val_seq,
                    y_val_seq_dict=self.y_val_seq_dict,
                    X_test_seq=self.X_test_seq,
                    y_test_seq_dict=self.y_test_seq_dict,
                    X_train_raw=self.X_train_raw,
                    X_val_raw=self.X_val_raw,
                    X_test_raw=self.X_test_raw,
                    y_train=self.df_train_original,
                    y_val=self.df_val_original,
                    y_test=self.df_test_original,
                    scaler=self.scaler, final_whitelist=self.final_whitelist, feature_engineer=self.feature_engineer,
                    train_indices=self.train_indices,
                    # <<< Pass original val/test indices >>>
                    val_indices=self.val_indices,
                    test_indices=self.test_indices,
                    io=self.io
                )
                
                # Handle potential re-processing: update sequences/scaler if new ones returned
                if scaler_returned is not None:
                    self.scaler = scaler_returned
                if X_train_seq_new is not None:
                    logger.info(f"Fold {self.current_fold}: Updating sequences due to GRU model loading/re-processing.")
                    self.X_train_seq = X_train_seq_new
                    self.y_train_seq_dict = y_train_seq_dict_new
                    self.train_indices = train_indices_new
                    self.X_val_seq = X_val_seq_new
                    self.y_val_seq_dict = y_val_seq_dict_new
                    self.val_indices = val_indices_new
                    self.X_test_seq = X_test_seq_new
                    self.y_test_seq_dict = y_test_seq_dict_new
                    self.test_indices = test_indices_new
                
                if self.gru_model is None:
                    raise SystemExit(f"Fold {self.current_fold}: Failed to train or load GRU model.")

                # k. Calibrate Probabilities (using stage function)
                # Replace method call with stage function call
                logger.info(f"Fold {self.current_fold}: Calibrating probabilities...")
                (self.optimal_T,
                 self.vector_cal_params,
                 self.optimized_edge_threshold,
                 p_cal_val_for_check, # Store intermediate result for validation checks
                 y_dir_val_for_check) = calibrate_probabilities_fold(
                     config=self.config,
                     current_fold=self.current_fold,
                     gru_model=self.gru_model,
                     gru_handler=self.gru_handler,
                     X_val_seq=self.X_val_seq,
                     y_val_seq_dict=self.y_val_seq_dict,
                     use_ternary=self.use_ternary,
                     calibrator=self.calibrator,
                     vector_calibrator=self.vector_calibrator,
                     fold_dirs=self.fold_dirs,
                     current_run_models_dir=self.current_run_models_dir,
                     run_id=self.run_id,
                     io=self.io
                 )
                 
                # Store results needed for validation checks on self
                self.p_cal_val_for_check = p_cal_val_for_check
                self.y_dir_val_for_check = y_dir_val_for_check
                # Optimized edge is already stored in self.optimized_edge_threshold by the call above

                # j.2. Perform GRU Validation Checks (Direct stage function call)
                # --- Capture report and check status --- #
                # Replace wrapper call with direct stage function call
                gru_val_report = run_gru_validation_checks_fold(
                    config=self.config,
                    current_fold=self.current_fold,
                    p_cal_val=self.p_cal_val_for_check, # Use intermediate result
                    y_dir_val=self.y_dir_val_for_check, # Use intermediate result
                    optimized_edge_threshold=self.optimized_edge_threshold,
                    use_ternary=self.use_ternary,
                    io=self.io
                )
                self.current_fold_report_data.append(gru_val_report)
                if gru_val_report.get('status') in ['failed', 'error']:
                    raise SystemExit(f"Fold {self.current_fold} failed GRU validation checks: {gru_val_report.get('reason')}")
                # --- End report handling --- #

                # k. Train/Load SAC Agent (Direct stage function call)
                logger.info(f"Fold {self.current_fold}: Proceeding to SAC training/loading step.")
                # Replace wrapper call with direct stage function call
                # Instantiate SACTrainer here if training
                sac_trainer_instance = None
                if self.config.get('sac', {}).get('train_sac', False):
                     if not self.io:
                          logger.error(f"Fold {self.current_fold}: IOManager is required to initialize SACTrainer for training. Skipping SAC training.")
                     else:
                          sac_trainer_instance = SACTrainer(config=self.config, io_manager=self.io) 
                          # Store instance if needed elsewhere? For now, just pass to stage func.
                          # self.sac_trainer = sac_trainer_instance 
                          
                self.sac_agent_load_path = train_or_load_sac_fold(
                    config=self.config,
                    current_fold=self.current_fold,
                    gru_model_run_id_loaded_from=self.gru_model_run_id_loaded_from,
                    base_models_dir_path=self.base_models_dir_path,
                    sac_trainer=sac_trainer_instance, # Pass instance or None
                    io=self.io
                )
                # Optional: Add a report entry for SAC training status?
                sac_train_status = {
                    'gate_name': 'sac_training',
                    'timestamp': datetime.now(timezone.utc).isoformat(),
                    'status': 'completed' if self.sac_agent_load_path else 'failed_or_skipped',
                    'agent_path': self.sac_agent_load_path,
                    'reason': None if self.sac_agent_load_path else 'SAC training/loading did not yield a valid agent path'
                }
                self.current_fold_report_data.append(sac_train_status)

                # l. Run Backtest (Direct stage function call)
                # --- Capture report and check status --- #
                # Replace wrapper call with direct stage function call
                # --- Prerequisites Check (simplified from wrapper) --- #
                backtest_prereqs_met = True
                if not all([
                    self.X_test_seq is not None, 
                    self.y_test_seq_dict is not None, 
                    self.test_indices is not None, 
                    self.df_test_original is not None
                ]):
                     logger.error(f"Fold {self.current_fold}: Missing test sequences, indices, or original test df. Cannot run backtest stage.")
                     backtest_prereqs_met = False
                     # Create an error report manually if prereqs fail
                     backtest_report = {
                         "gate_name": "backtest_stage_execution",
                         "status": "error",
                         "timestamp": datetime.now(timezone.utc).isoformat(),
                         "reason": "Missing required data for backtest stage.",
                         "metrics": {},
                         "thresholds": {}
                     }
                
                # --- Get Inputs for Backtest (simplified from wrapper) --- #
                p_raw_test_input: Optional[np.ndarray] = None 
                logits_test_input: Optional[np.ndarray] = None 
                if backtest_prereqs_met:
                    cal_method = self.config.get('calibration', {}).get('method', 'temperature')
                    if cal_method == 'vector':
                         if self.gru_handler:
                             logits_test_input = self.gru_handler.predict_logits(self.X_test_seq)
                             if logits_test_input is None: 
                                 logger.error(f"Fold {self.current_fold}: Failed to get logits for vector calibration in backtest.")
                                 backtest_prereqs_met = False
                                 backtest_report = { "status": "error", "reason": "Failed to get logits for backtest vector calibration." }
                         else:
                             logger.error(f"Fold {self.current_fold}: GRU handler not available for logits prediction.")
                             backtest_prereqs_met = False
                             backtest_report = { "status": "error", "reason": "GRU handler missing for logits prediction." }
                    else: # Temperature or other methods might need raw probabilities
                         if self.gru_handler:
                             predictions_dict = self.gru_handler.predict(self.X_test_seq)
                             if predictions_dict is None: 
                                 logger.error(f"Fold {self.current_fold}: Failed to get GRU predictions for backtest.")
                                 backtest_prereqs_met = False
                                 backtest_report = { "status": "error", "reason": "Failed to get GRU predictions for backtest." }
                             else:
                                 try:
                                     if self.use_ternary:
                                         p_raw_test_input = predictions_dict.get('dir3')
                                         if p_raw_test_input is None: raise KeyError("'dir3' missing")
                                     else:
                                         prob_key = next((k for k in predictions_dict if k != 'mu' and k != 'dir3_logits'), None)
                                         if prob_key:
                                             p_raw_test_input = predictions_dict[prob_key]
                                             if p_raw_test_input is not None and p_raw_test_input.ndim > 1: p_raw_test_input = p_raw_test_input.flatten()
                                         else:
                                             raise KeyError("Binary probability key not found")
                                 except KeyError as e:
                                     logger.error(f"Fold {self.current_fold}: Error extracting probabilities from GRU prediction dict: {e}")
                                     backtest_prereqs_met = False
                                     backtest_report = { "status": "error", "reason": f"Error extracting probabilities: {e}" }
                         else:
                             logger.error(f"Fold {self.current_fold}: GRU handler not available for prediction.")
                             backtest_prereqs_met = False
                             backtest_report = { "status": "error", "reason": "GRU handler missing for prediction." }
                             
                # --- Call Backtest Stage Function --- #
                if backtest_prereqs_met:
                    try:
                        results_df, metrics_dict, metrics_log, backtest_gate_report = run_backtest_fold(
                            config=self.config,
                            io=self.io,
                            current_fold=self.current_fold,
                            fold_dirs=self.fold_dirs,
                            sac_agent_load_path=self.sac_agent_load_path,
                            X_test_seq=self.X_test_seq,
                            y_test_seq_dict=self.y_test_seq_dict,
                            test_indices=self.test_indices, # Pass test_indices
                            df_test_original=self.df_test_original,
                            gru_handler=self.gru_handler,
                            calibrator=self.calibrator,
                            vector_calibrator=self.vector_calibrator,
                            initial_optimal_T=getattr(self, 'optimal_T', None),
                            initial_vector_params=getattr(self, 'vector_cal_params', None),
                            optimized_edge_threshold=self.optimized_edge_threshold,
                            p_raw_test=p_raw_test_input,
                            logits_test=logits_test_input,
                            use_ternary=self.use_ternary
                        )
                        # Store results on self if successful execution
                        self.backtest_results_df = results_df
                        self.backtest_metrics = metrics_dict
                        self.metrics_log_df = metrics_log
                        backtest_report = backtest_gate_report # Use the report from the stage func
                        logger.info(f"Fold {self.current_fold}: Backtest stage execution completed.")
                    except Exception as e:
                        logger.error(f"Fold {self.current_fold}: An unexpected error occurred calling run_backtest_fold stage: {e}", exc_info=True)
                        backtest_report = {
                            "gate_name": "backtest_stage_execution",
                            "status": "error",
                            "timestamp": datetime.now(timezone.utc).isoformat(),
                            "reason": f"Unhandled exception in run_backtest_fold: {e}",
                            "metrics": {},
                            "thresholds": {}
                        }
                # else: backtest_report was already set due to prereq failure
                
                # Add the final backtest report (either success/fail from gate or error report)
                self.current_fold_report_data.append(backtest_report)
                if backtest_report.get('status') == 'error': # Only raise SystemExit on execution error
                    raise SystemExit(f"Fold {self.current_fold} encountered an error during backtest stage: {backtest_report.get('reason')}")
                elif backtest_report.get('status') == 'failed': # Log warning if performance gate failed
                    logger.warning(f"Fold {self.current_fold} FAILED backtest performance gates: {backtest_report.get('reason')}. Proceeding to store metrics.")
                # --- End report handling --- #

                # m. Persist Fold Artefacts & Store Metrics
                logger.info(f"Storing metrics for Fold {self.current_fold}")
                if self.backtest_metrics is not None:
                     fold_metrics = self.backtest_metrics.copy()
                     fold_metrics['fold_number'] = self.current_fold
                     fold_metrics['train_start'] = train_start.isoformat() if train_start else None
                     fold_metrics['train_end'] = train_end.isoformat() if train_end else None
                     fold_metrics['val_start'] = val_start.isoformat() if val_start else None
                     fold_metrics['val_end'] = val_end.isoformat() if val_end else None
                     fold_metrics['test_start'] = test_start.isoformat() if test_start else None
                     fold_metrics['test_end'] = test_end.isoformat() if test_end else None
                     # Determine status based on metrics
                     if 'Annualized Sharpe Ratio' in fold_metrics and not pd.isna(fold_metrics['Annualized Sharpe Ratio']):
                         fold_metrics['status'] = 'success'
                     else:
                         fold_metrics['status'] = 'failed_backtest'
                     self.all_fold_metrics.append(fold_metrics)
                else:
                     logger.warning(f"Fold {self.current_fold}: No backtest metrics generated to store.")
                     self.all_fold_metrics.append({'fold_number': self.current_fold, 'status': 'failed_backtest', 'error': 'No metrics returned'})

                # Store SAC Agent Path if Trained Successfully for this fold
                if self.config.get('control', {}).get('train_sac', False) and self.sac_agent_load_path:
                     # Assuming sac_agent_load_path points to the *newly* trained agent if training occurred
                     if hasattr(self, 'sac_trainer') and self.sac_trainer and self.sac_trainer.last_saved_agent_path == self.sac_agent_load_path:
                          if os.path.exists(self.sac_agent_load_path):
                               logger.info(f"Fold {self.current_fold}: Storing successfully trained SAC agent path for aggregation: {self.sac_agent_load_path}")
                               all_successful_sac_agent_paths.append(self.sac_agent_load_path)
                          else:
                               logger.warning(f"Fold {self.current_fold}: SAC training reported success, but path {self.sac_agent_load_path} not found.")
                     elif self.sac_agent_load_path: # If path exists but wasn't from training this fold
                          pass # Don't add loaded agents to the aggregation list unless explicitly intended

                # <<< ADDED: Forward Baseline Check >>>
                # Derive y_test_ordinal if possible
                y_test_ordinal = None
                if self.y_test is not None and self.target_dir_col in self.y_test.columns:
                    if self.use_ternary:
                        # Convert one-hot list back to ordinal for test set
                        test_dir_raw = self.y_test[self.target_dir_col]
                        valid_mask_test = test_dir_raw.notna() & test_dir_raw.apply(lambda x: isinstance(x, list) and len(x) == 3)
                        if valid_mask_test.any():
                             ordinal_values_test = test_dir_raw[valid_mask_test].apply(np.argmax)
                             y_test_ordinal = pd.Series(np.nan, index=test_dir_raw.index)
                             y_test_ordinal[valid_mask_test] = ordinal_values_test
                             logger.info(f"Fold {self.current_fold}: Extracted ordinal test labels for forward check. Count: {valid_mask_test.sum()}")
                        else:
                             logger.warning(f"Fold {self.current_fold}: No valid ternary labels found in y_test for forward check.")
                             y_test_ordinal = pd.Series(dtype=float) # Empty series
                    else:
                        # Binary case: convert float (0.0/1.0) to int
                        y_test_ordinal = self.y_test[self.target_dir_col].astype(int)
                else:
                     logger.warning(f"Fold {self.current_fold}: Cannot run forward baseline check: y_test or target_dir_col ('{self.target_dir_col}') is missing.")

                if y_test_ordinal is not None:
                    forward_check_passed = self.baseline_checker.run_forward_baseline_check(
                        X_train_fold=self.X_train_scaled,
                        y_train_fold_ordinal=self.y_dir_train_ordinal, 
                        X_test_fold=self.X_test_scaled,
                        y_test_fold_ordinal=y_test_ordinal, 
                        fold_num=self.current_fold,
                        io=self.io
                    )
                    if not forward_check_passed:
                        # Raise SystemExit to be caught by the except block below
                        raise SystemExit(f"Fold {self.current_fold}: FAILED Forward Baseline Check Gate")
                # <<< END ADDED >>>


            except SystemExit as fold_exit:
                 # Log the reason for stopping the fold and continue
                 logger.error(f"Fold {self.current_fold} processing stopped: {fold_exit}. Skipping subsequent stages (including GRU Validation & SAC) for this fold.") 
                 # Store failure reason more specifically
                 self.all_fold_metrics.append({'fold_number': self.current_fold, 'status': f'failed_{fold_exit}', 'error': str(fold_exit)}) # <<< Updated status logging
                 # Log that SAC wasn't reached
                 logger.info(f"Fold {self.current_fold}: *** Did NOT reach SAC training step due to SystemExit: {fold_exit}. ***")
                 continue # Proceed to the next fold
            except Exception as e: # Catch any exception during fold processing
                 # Log the unexpected error and continue
                 logger.error(f"Fold {self.current_fold}: Unhandled exception occurred: {e}. Skipping to next fold.", exc_info=True) # Add exc_info for traceback
                 self.all_fold_metrics.append({'fold_number': self.current_fold, 'status': 'error', 'error': str(e)}) # <<< Updated status logging
                 # Log that SAC wasn't reached
                 logger.info(f"Fold {self.current_fold}: *** Did NOT reach SAC training step due to unhandled exception. ***")
                 continue # Proceed to the next fold

            # --- End Fold Loop --- #
            logger.info(f"=== Finished Processing Fold {self.current_fold} ===")
            # If only single split, break after first iteration - Generator handles this now
            # No explicit break needed here anymore.

        # 4. Aggregate Fold Metrics & Final Decision
        release_decision_passed = False
        if self.all_fold_metrics:
            self.aggregated_metrics = self.aggregate_fold_metrics(self.all_fold_metrics)
            logger.info("--- Aggregated Walk-Forward Metrics --- ")
            # Use json dumps for pretty printing dict/nested dict
            # <<< ADD CHECK FOR NON-NONE AGGREGATED METRICS BEFORE LOGGING/SAVING >>>
            if self.aggregated_metrics is not None:
                logger.info(json.dumps(self.aggregated_metrics, default=lambda x: str(x) if isinstance(x, (pd.Timestamp, np.ndarray)) else x, indent=2))
                if self.io:
                     self.io.save_json(self.aggregated_metrics, 'aggregated_wf_metrics', section='results')
            else:
                logger.error("Aggregation resulted in None. Cannot log or save aggregated metrics.")
            # <<< END CHECK >>>

            # <<< ADD CHECK FOR NON-NONE AGGREGATED METRICS BEFORE FINAL DECISION >>>
            if self.aggregated_metrics is not None:
                release_decision_passed = self.final_release_decision(self.aggregated_metrics)
            else:
                 logger.error("Skipping final release decision because aggregated metrics are None.")
        else:
            logger.warning("No fold metrics were generated. Skipping aggregation and final decision.")
            self.aggregated_metrics = {} # Ensure defined as empty dict if no metrics

        # Log Final Status
        if release_decision_passed:
             logger.info(f"--- Pipeline Run {self.run_id} finished successfully and meets release criteria. ---")
        else:
             logger.error(f"--- Pipeline Run {self.run_id} finished but FAILED to meet release criteria. See aggregated metrics and logs. ---")

        # 5. Aggregate SAC Agents (Optional)
        if self.config.get('sac_aggregation', {}).get('enabled', False):
             if all_successful_sac_agent_paths:
                  self.aggregate_sac_agents(all_successful_sac_agent_paths)
             else:
                  logger.warning("SAC agent aggregation enabled, but no successfully trained fold agents were found/stored. Skipping aggregation.")
        else:
             logger.info("SAC agent aggregation disabled. Skipping.")


        # --- Log Summary (using aggregated_metrics dict) --- #
        logger.info(f"--- Aggregated Fold Metrics Summary ---")
        agg_status = self.aggregated_metrics.get("aggregation_status", "unknown")
        total_folds_agg = self.aggregated_metrics.get("total_folds", 0)
        fully_successful_agg = self.aggregated_metrics.get("fully_successful_folds", 0)
        partially_completed_agg = self.aggregated_metrics.get("partially_completed_folds", 0)
        failed_early_agg = self.aggregated_metrics.get("failed_early_folds", 0)
        pass_rate_agg = self.aggregated_metrics.get("pass_rate", 0.0)
        failed_details_agg = self.aggregated_metrics.get("failed_fold_details", {})

        logger.info(f"Aggregation Status: {agg_status}")
        logger.info(f"Total Folds Attempted: {total_folds_agg}")
        logger.info(f"  - Fully Successful (passed all gates): {fully_successful_agg}")
        logger.info(f"  - Partially Completed (failed late gate): {partially_completed_agg}")
        logger.info(f"  - Failed Early / Error: {failed_early_agg}")
        logger.info(f"Overall Pass Rate (Fully Successful / Total): {pass_rate_agg:.2%}")
        if failed_details_agg:
            logger.info("Details for Non-Successful Folds:")
            # Sort by fold number if possible
            try:
                 sorted_failed_details = sorted(failed_details_agg.items(), key=lambda item: int(item[0]))
            except ValueError:
                 sorted_failed_details = failed_details_agg.items() # Fallback if fold numbers aren't ints
                 
            for fold, details in sorted_failed_details:
                status = details.get('status', 'unknown')
                error = details.get('error', '')
                logger.info(f"  - Fold {fold}: Status='{status}', Error='{error}'")
        # --- End Log Summary --- #

    def aggregate_fold_metrics(self, all_fold_metrics: List[dict]) -> dict:
        """
        Aggregates metrics collected from successful walk-forward folds.

        Args:
            all_fold_metrics (List[dict]): A list where each element is a dictionary
                                          of metrics from a single fold, potentially
                                          including 'status' or 'error' keys for failed folds.

        Returns:
            dict: A dictionary containing summary statistics (mean, std, min, max, count)
                  for key performance metrics across fully successful folds, plus detailed fold counts and pass rate.
                  Returns an error dict if no folds completed sufficiently for aggregation.
        """
        num_total_folds = len(all_fold_metrics)
        if num_total_folds == 0:
            logger.warning("No fold metrics provided for aggregation.")
            return {"aggregation_status": "failed", "reason": "No fold metrics"}

        # --- Categorize Folds ---
        fully_successful_folds_data = []
        partially_completed_folds_data = []
        failed_early_folds_data = []
        failed_fold_details = {} # Store status/error for non-successful folds

        # Define statuses indicating partial completion (ran significant steps but failed a late gate)
        # Use startswith matching for flexibility (e.g., 'failed_backtest', 'failed_gru_gate', 'failed_forward_baseline')
        partial_completion_prefixes = ('failed_backtest', 'failed_gru_gate', 'failed_forward_baseline')
        # Define statuses indicating full success
        full_success_status = 'success'

        for m in all_fold_metrics:
            fold_num = m.get('fold_number', 'unknown')
            status = m.get('status', 'unknown')
            error_msg = m.get('error', '')

            # Ensure status is a string for startswith check
            status_str = str(status) if status is not None else 'unknown'

            # Check for full success first
            if status_str == full_success_status and 'Annualized Sharpe Ratio' in m and not pd.isna(m['Annualized Sharpe Ratio']):
                fully_successful_folds_data.append(m)
            else:
                # Store details for all non-successful folds
                failed_fold_details[fold_num] = {'status': status_str, 'error': str(error_msg)}
                # Categorize further
                is_partial = False
                for partial_prefix in partial_completion_prefixes:
                    if status_str.startswith(partial_prefix):
                        partially_completed_folds_data.append(m)
                        is_partial = True
                        break
                if not is_partial:
                    failed_early_folds_data.append(m) # Assume others failed earlier or had generic error

        num_fully_successful = len(fully_successful_folds_data)
        num_partially_completed = len(partially_completed_folds_data)
        num_failed_early = len(failed_early_folds_data)
        # Verification: num_total_folds should equal sum of categories
        if num_total_folds != (num_fully_successful + num_partially_completed + num_failed_early):
            logger.warning(f"Fold count mismatch during aggregation: Total={num_total_folds}, Categories Sum = {num_fully_successful + num_partially_completed + num_failed_early}. Check status strings/logic.")

        pass_rate = num_fully_successful / num_total_folds if num_total_folds > 0 else 0.0
        # --- Log Summary within aggregation function - can be removed if logged outside ---
        # logger.info(f"--- Aggregating Fold Metrics --- ")
        # logger.info(f"Total Folds Attempted: {num_total_folds}")
        # ... (rest of logging removed as it's now done outside) ...
        # --- End Log Summary ---

        # ... (rest of the code remains unchanged)

    def final_release_decision(self, aggregated_metrics: dict) -> bool:
        """
        Makes a final decision based on aggregated walk-forward metrics.

        Args:
            aggregated_metrics (dict): Dictionary of aggregated metrics from
                                        aggregate_fold_metrics. Can be None if aggregation failed.

        Returns:
            bool: True if the performance meets release criteria, False otherwise.
        """
        logger.info("--- Making Final Release Decision based on Aggregated Metrics ---")

        # <<< Updated check for aggregation status (Handles None input) >>>
        if aggregated_metrics is None:
            logger.error("Final Release Decision: FAILED - Aggregated metrics object is None. Cannot evaluate release criteria.")
            return False
            
        if aggregated_metrics.get('aggregation_status') not in ['success', 'no_numeric_metrics', 'no_successful_folds']:
            logger.error(f"Final Release Decision: FAILED - Aggregation status indicates failure ('{aggregated_metrics.get('aggregation_status', 'N/A')}'). Cannot evaluate release criteria.")
            return False

        # --- Load Release Criteria from Config --- #
        gate_config = self.config.get('validation_gates', {}).get('final_release', {})
        criteria = {
            'min_successful_folds_pct': gate_config.get('min_successful_folds_pct', 0.75),
            'median_sharpe_threshold': gate_config.get('median_sharpe_threshold', 1.3),
            # Add other criteria as needed (e.g., max_drawdown_max_threshold)
            'max_drawdown_max_threshold': gate_config.get('max_drawdown_max_threshold', 20.0) # Example
        }
        logger.info(f"Using final release criteria from config: {criteria}")
        # --- End Load Criteria --- #

        # --- Check Criteria --- #
        passed_all = True
        reasons = []

        # 1. Successful Folds Percentage (Pass Rate)
        total_folds = aggregated_metrics.get('total_folds', 0)
        # Use the calculated pass_rate from the aggregation dict
        pass_rate = aggregated_metrics.get('pass_rate', 0.0)

        if total_folds > 0:
            if pass_rate < criteria['min_successful_folds_pct']:
                passed_all = False
                reasons.append(f"Insufficient pass rate (fully successful folds): {pass_rate:.1%} < {criteria['min_successful_folds_pct']:.1%}")
            else:
                reasons.append(f"Pass rate OK: {pass_rate:.1%} >= {criteria['min_successful_folds_pct']:.1%}")
        else:
            passed_all = False # Should have been caught by aggregation status check, but belt-and-suspenders
            reasons.append("No folds were run.")

        # --- Check Performance Metrics (only if aggregation status allows) ---
        if aggregated_metrics.get('aggregation_status') == 'success':
            perf_metrics = aggregated_metrics.get('aggregated_performance_metrics', {})
            sharpe_key = None
            if "Annualized Sharpe Ratio (Re-centred)" in perf_metrics:
                sharpe_key = "Annualized Sharpe Ratio (Re-centred)"
            elif "Annualized Sharpe Ratio" in perf_metrics:
                sharpe_key = "Annualized Sharpe Ratio"

            # 2. Median Sharpe Ratio
            if sharpe_key:
                if hasattr(self, 'all_fold_metrics') and self.all_fold_metrics:
                    successful_fold_metrics_list = [m for m in self.all_fold_metrics if m.get('status') == 'success' and sharpe_key in m and not pd.isna(m[sharpe_key])]
                    if successful_fold_metrics_list:
                        successful_fold_metrics_df = pd.DataFrame(successful_fold_metrics_list)
                        median_sharpe = successful_fold_metrics_df[sharpe_key].median()
                        if np.isnan(median_sharpe) or median_sharpe < criteria['median_sharpe_threshold']:
                            passed_all = False
                            reasons.append(f"Median Sharpe ({sharpe_key}) too low: {median_sharpe:.2f} < {criteria['median_sharpe_threshold']}")
                        else:
                            reasons.append(f"Median Sharpe OK: {median_sharpe:.2f} >= {criteria['median_sharpe_threshold']}")
                    else:
                        passed_all = False
                        reasons.append(f"No fully successful folds with valid Sharpe metric '{sharpe_key}' found for median calculation.")
                else:
                    passed_all = False
                    reasons.append("Could not access per-fold metrics to calculate median Sharpe.")
            else:
                # Only fail if the sharpe threshold is actually set in criteria
                if 'median_sharpe_threshold' in criteria:
                    passed_all = False
                    reasons.append("Sharpe metric required by criteria but not found in aggregation results for median check.")
                else:
                    reasons.append("Sharpe metric not found in aggregation results (median check skipped as not required by criteria).")


            # 3. Max Drawdown
            max_dd_key = "Max Drawdown (%)"
            if 'max_drawdown_max_threshold' in criteria:
                 if max_dd_key in perf_metrics and 'max' in perf_metrics[max_dd_key]:
                    max_dd_max = perf_metrics[max_dd_key].get('max', np.nan)
                    max_dd_threshold = criteria['max_drawdown_max_threshold']
                    if np.isnan(max_dd_max):
                         passed_all = False # Treat NaN as failure if threshold is set
                         reasons.append(f"Max Drawdown is NaN, cannot compare to threshold {max_dd_threshold}%.")
                    elif max_dd_max > max_dd_threshold:
                        passed_all = False
                        reasons.append(f"Max Drawdown exceeded limit: {max_dd_max:.2f}% > {max_dd_threshold}%")
                    else:
                        reasons.append(f"Max Drawdown OK: {max_dd_max:.2f}% <= {max_dd_threshold}%")
                 else:
                    passed_all = False # Metric required but not found or 'max' stat missing
                    reasons.append(f"Required Max Drawdown metric ('{max_dd_key}' with 'max' stat) not found in aggregation.")
            # ... Add similar checks for other criteria ...

        elif aggregated_metrics.get('aggregation_status') == 'no_numeric_metrics':
            reasons.append("Cannot check performance metrics (Sharpe, Drawdown) as none were aggregated.")
            # Decision depends on whether min_successful_folds_pct was met
            if not passed_all: # If pass rate check already failed
                 pass # Keep passed_all as False
            else: # Pass rate was ok, but no metrics to check
                 # Decide if this constitutes a failure - let's say yes if performance metrics are crucial
                 if 'median_sharpe_threshold' in criteria or 'max_drawdown_max_threshold' in criteria:
                      passed_all = False
                      reasons.append("Release criteria require performance metrics, but none were available for checking.")

        elif aggregated_metrics.get('aggregation_status') == 'no_successful_folds':
             reasons.append("Cannot check performance metrics (Sharpe, Drawdown) as no folds were fully successful.")
             # pass_rate check should have already failed if min_successful_folds_pct > 0
             passed_all = False # Explicitly set to False

        # --- Log Final Decision --- #
        decision_str = "PASSED" if passed_all else "FAILED"
        logger.info(f"Final Release Decision: {decision_str}")
        for reason in reasons:
            logger.info(f"  - {reason}")

        return passed_all
        # Removed dangling else/except block that didn't belong here

        # ... (rest of the code remains unchanged)

    def aggregate_sac_agents(self, agent_paths: List[str]):
        """Averages the weights of multiple SAC agents saved during walk-forward folds."""
        if not agent_paths:
            logger.warning("No SAC agent paths provided for aggregation. Skipping.")
            return

        logger.info(f"Starting aggregation of {len(agent_paths)} SAC agents...")
        
        # --- Load State Dicts --- #
        all_state_dicts = []
        reference_state_dict = None
        for path in agent_paths:
            try:
                # Load onto CPU to avoid potential GPU memory issues during aggregation
                state_dict = torch.load(path, map_location=torch.device('cpu'))
                
                # Basic validation: check if it's a dictionary
                if not isinstance(state_dict, dict):
                     logger.warning(f"Skipping agent at {path}: Loaded object is not a dictionary.")
                     continue
                     
                # Further validation: Check if essential keys exist (adjust based on SACTradingAgent structure)
                # Example check (modify based on actual keys):
                required_keys = ['actor_state_dict', 'critic_state_dict'] 
                if not all(key in state_dict for key in required_keys):
                     logger.warning(f"Skipping agent at {path}: Missing required keys ({required_keys}).")
                     continue
                     
                all_state_dicts.append(state_dict)
                if reference_state_dict is None:
                     reference_state_dict = state_dict # Use first valid one as structure reference
                     
                # Optional: Add more detailed structure comparison here if needed
                # e.g., check if keys within actor/critic state dicts match reference
                     
            except FileNotFoundError:
                logger.warning(f"Skipping agent aggregation: File not found at {path}")
            except Exception as e:
                logger.warning(f"Skipping agent aggregation: Failed to load agent from {path}: {e}")

        if len(all_state_dicts) < 2:
            logger.warning(f"Need at least two valid agents for aggregation, found {len(all_state_dicts)}. Skipping aggregation.")
            return
            
        if reference_state_dict is None:
             logger.error("Could not load any valid reference state dict. Cannot proceed with aggregation.")
             return
             
        logger.info(f"Successfully loaded {len(all_state_dicts)} valid state dictionaries for aggregation.")

        # --- Average Parameters --- #
        # Deepcopy the structure from the reference
        aggregated_state_dict = copy.deepcopy(reference_state_dict)
        
        # Iterate through components (actor, critic, etc.)
        for component_key in reference_state_dict.keys(): # e.g., 'actor_state_dict', 'critic_state_dict'
             if isinstance(reference_state_dict[component_key], OrderedDict):
                  logger.debug(f"Aggregating parameters for component: {component_key}")
                  # Iterate through parameters within the component's state dict
                  for param_key in reference_state_dict[component_key].keys():
                       # Sum tensors from all loaded state dicts
                       summed_tensor = None
                       valid_tensors_count = 0
                       for sd in all_state_dicts:
                           # Check if the component and parameter exist in this state dict
                           if component_key in sd and isinstance(sd[component_key], OrderedDict) and param_key in sd[component_key]:
                                tensor = sd[component_key][param_key]
                                if summed_tensor is None:
                                    summed_tensor = tensor.clone().detach().float() # Use float for averaging
                                else:
                                     # Ensure tensors are compatible before adding
                                     if summed_tensor.shape == tensor.shape:
                                          summed_tensor += tensor.float()
                                     else:
                                          logger.warning(f"Shape mismatch for {component_key}.{param_key}. Skipping tensor from one agent.")
                                          continue # Skip this tensor, but potentially count others?
                                valid_tensors_count += 1
                           else:
                                logger.warning(f"Parameter {component_key}.{param_key} not found in one of the state dicts. Skipping.")

                       # Calculate average if tensors were found and summed
                       if summed_tensor is not None and valid_tensors_count > 0:
                            averaged_tensor = summed_tensor / valid_tensors_count
                            # Update the aggregated state dict
                            aggregated_state_dict[component_key][param_key] = averaged_tensor
                            logger.debug(f"  Averaged {component_key}.{param_key} from {valid_tensors_count} agents.")
                       elif valid_tensors_count == 0:
                            logger.warning(f"  Could not find any valid tensors for {component_key}.{param_key}. Keeping reference value.")
             else:
                  # Handle non-state-dict items if necessary (e.g., config, optimizers - usually not averaged)
                  logger.debug(f"Skipping non-state-dict component: {component_key}")
                  # Keep the value from the reference dict for non-parameter items
                  aggregated_state_dict[component_key] = reference_state_dict[component_key]

        # --- Save Aggregated Agent --- #
        if self.io:
            try:
                agg_agent_filename = 'sac_agent_aggregated' # Consistent filename
                save_path = self.io.path('models', agg_agent_filename, suffix='.pth') # Save in main run models dir
                torch.save(aggregated_state_dict, save_path)
                logger.info(f"Successfully saved aggregated SAC agent state dictionary to: {save_path}")
            except Exception as e:
                logger.error(f"Failed to save aggregated SAC agent: {e}", exc_info=True)
        else:
            logger.warning("IOManager not available. Cannot save aggregated SAC agent.")
       