feat: Add imputed transition logging to SACTrainer and update docs

gru_sac_predictor/README.md

@@ -1,6 +1,6 @@
-# GRU + SAC Crypto Trading System (v3 - Consolidated & Enhanced)
+# GRU + SAC Crypto Trading System (v3 - Refactored & Enhanced)
 
-This project implements a cryptocurrency trading system using a GRU model for market prediction and a Soft Actor-Critic (SAC) agent for position sizing. This version reflects significant refactoring, feature additions, and the consolidation of model logic.
+This project implements a cryptocurrency trading system using a GRU model for market prediction and a Soft Actor-Critic (SAC) agent for position sizing. This version reflects significant refactoring for modularity, feature additions, and the consolidation of model logic.
 
 The core idea is to decouple prediction and action:
 1.  A **GRU model** (v2 or v3 architecture, selected via config) forecasts future log-returns (μ̂) and class probabilities (binary p(up) or ternary p(down, flat, up)).
@@ -11,32 +11,71 @@ This approach aims for a robust system where the RL agent focuses solely on risk
 
 ## Key Features & Enhancements
 
-*   **Consolidated GRU Logic:** Both v2 and v3 GRU model architectures are now implemented and managed within `src/gru_model_handler.py`.
-*   **Walk-Forward Validation:** Replaces static train/val/test splits with a robust walk-forward validation framework (`TradingPipeline.execute`, `_generate_walk_forward_folds`) for more realistic performance estimation.
-*   **Hyperparameter Optimization (Optuna):** Integrated Optuna sweep for GRU hyperparameters (`src/gru_hyper_tuner.py`) with restricted search space, configurable objective (`edge_acc - brier`), and Keras callback for efficient pruning based on `val_loss`.
+*   **Modular Pipeline Structure:** Pipeline logic is refactored into stage-specific functions within `src/pipeline_stages/` for improved readability, maintainability, and testability (see Project Structure).
+*   **Consolidated GRU Logic:** Both v2 and v3 GRU model architectures are implemented and managed within `src/gru_model_handler.py`.
+*   **Walk-Forward Validation:** Robust walk-forward validation framework (`TradingPipeline.execute`, `_generate_walk_forward_folds`) for realistic performance estimation.
+*   **Hyperparameter Optimization (Optuna):** Integrated Optuna sweep for GRU hyperparameters (`src/gru_hyper_tuner.py`) with configurable search space, objective, and pruning.
 *   **Advanced Calibration:**
     *   Supports Temperature and Vector Scaling (`calibration.method`) with optional L2 regularization (`calibration.l2_lambda`).
-    *   Optimizes edge threshold via Youden's J on validation data (`calibration.optimize_edge_threshold`). Saved per fold (`optimized_edge_threshold_fold_N.txt`) and used consistently.
-    *   **Rolling Calibration (Experimental):** Implemented within `Backtester` to refit the calibrator periodically during the backtest (`calibration.rolling_enabled`, `recalibrate_every_n`, `recalibration_window`). Uses the **static** calibration from training time if SAC training is active to prevent lookahead.
+    *   Optimizes edge threshold via Youden's J on validation data (`calibration.optimize_edge_threshold`).
+    *   **Rolling Calibration (Experimental):** Implemented within `Backtester` to refit the calibrator periodically during the backtest (`calibration.rolling_enabled`, `recalibrate_every_n`, `recalibration_window`).
     *   **Coverage Alarm (ECE-based):** Optional alarm triggers early recalibration if **Expected Calibration Error (ECE)** exceeds a threshold (`calibration.coverage_alarm_enabled`, `ece_recalibration_threshold`).
-*   **Prioritized Experience Replay (PER):** Implemented for SAC training (`sac.use_per`) with **TD-error clipping** and **alpha annealing** (linear decay). Logs TD error distribution statistics.
-*   **SAC Enhancements:** Reward scaling, state normalization (`MeanStdFilter`), configurable **action penalty** (default: `0.01 / transaction_cost`), oracle seeding with **Importance Sampling (IS) weight decay** (`per_seed_decay_steps`).
-*   **Refined Validation Gates:** Configurable thresholds (`validation_gates`) for:
-    *   **Baseline Gate:** Checks Logistic Regression CI on **raw/engineered** training features before scaling.
-    *   **GRU Gate:** Checks Edge Acc CI and Brier score on validation set after calibration, using the fold's **determined edge threshold**.
-    *   **Final Release Decision:** Checks aggregated metrics (e.g., **median Sharpe ≥ 1.3**, **≥ 75% successful folds**) across all successful folds. Backtest gate failures *per fold* are logged but do not halt the entire pipeline.
-*   **Micro-structure Features:** Added bar-level features (`FeatureEngineer._add_microstructure_features`) with **NaN guards** for robustness.
-*   **Leakage Guard:** Feature calculations use `shift(1)`. Selection occurs on **raw/engineered features** and includes correlation check (`corr(ret+h, feat_t-1)`) against future returns. Minimal whitelist applied before VIF.
-*   **Configuration:** Centralized and expanded `config.yaml` with annotations for mutually exclusive options (e.g., walk-forward vs static split).
+*   **Prioritized Experience Replay (PER):** Implemented for SAC training (`sac.use_per`) with TD-error clipping and alpha annealing.
+*   **SAC Enhancements:** Reward scaling, state normalization (`MeanStdFilter`), configurable action penalty, oracle seeding with Importance Sampling (IS) weight decay.
+*   **Refined Validation Gates:** Configurable thresholds (`validation_gates`) for baseline checks, GRU validation, fold backtest performance, and final release decisions.
+*   **Micro-structure Features:** Added bar-level features (`FeatureEngineer._add_microstructure_features`) with NaN guards.
+*   **Leakage Guard:** Feature calculations use `shift(1)`. Selection includes correlation check against future returns. Minimal whitelist applied before VIF.
+*   **Configuration:** Centralized and expanded `config.yaml`.
 *   **Output Management:** Standardized output structure via `IOManager` and `LoggerSetup`.
-*   **SAC Agent Aggregation:** Optional post-processing step to average weights from agents trained across successful folds (`TradingPipeline.aggregate_sac_agents`, `sac_aggregation.enabled`).
+*   **SAC Agent Aggregation:** Optional post-processing step to average weights from agents trained across successful folds.
+
+## Data Quality
+
+This system includes mechanisms to handle potential missing data points (bars) in the input time series. This is crucial for maintaining data integrity and preventing errors during feature engineering and model training.
+
+**Handling Missing Bars:**
+
+*   **Detection:** The pipeline automatically detects missing bars based on the expected `data.bar_frequency` (e.g., "1T" for 1 minute) after initial data loading.
+*   **Reporting:** A warning is logged detailing the total number of missing bars found and the length of the longest consecutive gap. A summary report (`missing_bars_summary.json`) is saved in the run's results directory.
+*   **Filling Strategies:** Several strategies are available, configured via `data.missing.strategy`:
+    *   `"drop"`: No filling is performed. Missing bars remain gaps or NaNs. (Use with caution).
+    *   `"neutral"`: Forward-fills the 'close' price, sets 'open', 'high', 'low' equal to the filled 'close', and sets 'volume' to 0 for imputed bars.
+    *   `"ffill"`: Forward-fills all OHLCV columns, then back-fills any remaining NaNs at the beginning.
+    *   `"interpolate"`: Interpolates missing values using the method specified in `data.missing.interpolate.method` (e.g., 'linear') up to a `limit` defined in `data.missing.interpolate.limit`.
+*   **Imputed Flag:** After filling, a boolean column `bar_imputed` is added to the DataFrame, marking rows that were originally missing.
+*   **Max Gap Check:** The pipeline will raise an error if the longest detected consecutive gap exceeds `data.missing.max_gap`.
+
+**Impact on Downstream Components:**
+
+*   **Feature Engineering:** Features are calculated on the potentially gap-filled data.
+*   **Sequence Creation:** Sequences containing imputed bars can be optionally dropped before GRU training, controlled by `gru.drop_imputed_sequences`.
+*   **GRU Model:** The `bar_imputed` flag is included as a feature input to the GRU model, allowing it to potentially learn patterns related to imputed data.
+*   **SAC Environment:** The `TradingEnv` is aware of imputed bars. The behavior during an imputed step is controlled by `sac.imputed_handling`:
+    *   `"skip"`: The environment skips the step, no action is taken, no reward is given, and the transition is not added to the replay buffer.
+    *   `"hold"`: The agent's action is overridden to maintain its current position. The step proceeds normally otherwise (reward calculated based on held position).
+    *   `"penalty"`: The agent's chosen action is taken, but a penalty reward (based on `sac.action_penalty`) is applied instead of the normal PnL reward.
+
+**Recommended Defaults:**
+
+Using `"neutral"` or `"ffill"` for `strategy` is generally recommended for continuous time series. `max_gap` should be set to a reasonably small number (e.g., 5-10) to avoid filling excessively long gaps with potentially inaccurate data. For the SAC environment, `"hold"` or `"skip"` are common choices, depending on whether you want the agent to explicitly learn from imputed steps (or lack thereof).
 
 ## System Design & Workflow
 
 The system is orchestrated by `run.py`, which sets up logging and I/O via `LoggerSetup` and `IOManager`, then instantiates and executes the `TradingPipeline` class (`src/trading_pipeline.py`). The pipeline follows a sequence of steps, potentially looped for Walk-Forward validation.
 
+### Pipeline Stages (Refactored)
+
+The core logic for each step in the pipeline has been moved into dedicated functions within the `src/pipeline_stages/` directory. The `TradingPipeline` class now acts primarily as an orchestrator, calling these stage functions in sequence and managing the overall state and data flow.
+
+*   **`src/pipeline_stages/data_processing.py`**: Handles loading, initial preprocessing, feature engineering, labeling, and data splitting logic for each fold.
+*   **`src/pipeline_stages/feature_processing.py`**: Manages feature scaling, selection (L1+VIF), and pruning based on the selected whitelist.
+*   **`src/pipeline_stages/sequence_creation.py`**: Creates input sequences suitable for the GRU model from the processed feature data.
+*   **`src/pipeline_stages/modelling.py`**: Contains functions for training/loading the GRU model (including hyperparameter tuning), calibrating probabilities (Temperature/Vector Scaling, edge threshold optimization), training/loading the SAC agent, and aggregating SAC agents.
+*   **`src/pipeline_stages/evaluation.py`**: Includes functions for running baseline checks (Logistic Regression), performing GRU validation checks (Edge Accuracy, Brier Score), and executing the main backtest simulation (instantiating and running the `Backtester`).
+
+### Workflow Diagram
+
 ```mermaid
-%%{init: {'themeVariables': { 'fontSize': '26px' }}}%%
 graph TD
     A[run.py: Init Logger/IOManager/Config] --> B(TradingPipeline);
 
@@ -47,21 +86,21 @@ graph TD
         D --> F[Select Fold Data];
         E --> F;
 
-        subgraph Fold Processing [Fold Processing]
+        subgraph Fold Processing [Fold Processing - Calls Stage Functions]
             direction TB
-            F --> G[Engineer Features];
-            G --> H[Define Labels & Align];
-            H --> I[Split Fold Data];
-            I --> J1[Baseline Check];
-            J1 -- Pass --> L[Select Features];
-            L --> K[Scale Features];
-            K --> M[Prune Scaled Features];
-            M --> N[Create Sequences];
-            N --> O[Train/Load GRU];
-            O --> P[Calibrate Probabilities];
-            P --> R[GRU Validation Gate];
-            R -- Pass --> S[Train/Load SAC Agent];
-            S --> T[Run Backtest];
+            F --> G[data_processing: Engineer Features];
+            G --> H[data_processing: Define Labels & Align];
+            H --> I[data_processing: Split Fold Data];
+            I --> J1[evaluation: Baseline Check];
+            J1 -- Pass --> L[feature_processing: Select Features];
+            L --> K[feature_processing: Scale Features];
+            K --> M[feature_processing: Prune Scaled Features];
+            M --> N[sequence_creation: Create Sequences];
+            N --> O[modelling: Train/Load GRU];
+            O --> P[modelling: Calibrate Probabilities];
+            P --> R[evaluation: GRU Validation Gate];
+            R -- Pass --> S[modelling: Train/Load SAC Agent];
+            S --> T[evaluation: Run Backtest];
             T --> U[Record Fold Results];
         end
 
@@ -73,172 +112,144 @@ graph TD
 
     B --> Walk-ForwardLoop;
     X1 --> Y[Aggregate Fold Metrics];
-    Y --> Z[Aggregate SAC Agents];
+    Y --> Z[modelling: Aggregate SAC Agents];
     Z --> Z1[Final Release Decision];
     Z1 --> Z_End([End Pipeline Run]);
 
 ```
-*Diagram outlines the consolidated v3 pipeline flow after `revisions.txt` modifications.*
+*Diagram outlines the consolidated v3 pipeline flow, highlighting calls to stage functions.*
 
 ### Detailed Steps (Walk-Forward Enabled)
 
-1.  **Initialization (`run.py`):**
-    *   Parses args (`--config`, etc.).
-    *   Loads `config.yaml`.
-    *   Initializes `IOManager`, `LoggerSetup`.
-    *   Instantiates `TradingPipeline`.
-2.  **Data Loading (`TradingPipeline.load_and_preprocess_data`):**
-    *   Loads the *entire* raw dataset specified in the config.
-3.  **Fold Generation (`TradingPipeline._generate_walk_forward_folds`):**
-    *   Based on `walk_forward` config (train/val/test/step days), yields date ranges for each fold.
-4.  **Fold Loop (`TradingPipeline.execute`):** Iterates through generated folds.
-    *   **Select Fold Data:** Extracts raw data corresponding to the current fold's (Train+Val+Test) date range.
-    *   **Feature Engineering (`engineer_features`):** Computes base, TA, and micro-structure features (with NaN guards) on the fold's raw data (using `shift(1)` for time-dependent features).
-    *   **Labeling (`define_labels_and_align`):** Calculates forward returns and target labels (binary/ternary) for the fold's engineered data.
-    *   **Split Fold Data (`split_data`):** Splits the fold's labeled data into `train`, `val`, and `test` sets based on the fold's date ranges. Stores results like `self.X_train_raw`, `self.y_val`, `self.df_test_original`.
-    *   **Baseline Gate (`run_baseline_checks`):** Trains/validates Logistic Regression on fold's ***raw/engineered*** training features. Exits fold if CI lower bound < config threshold (`validation_gates.baseline`). Saves `baseline_report_fold_N.txt`.
-    *   **Select Features (`select_and_prune_features` - Selection Part):** Performs leakage check (`corr(ret+h, feat_t-1)`) and L1 + VIF selection (applying minimal whitelist *before* VIF) on fold's *raw/engineered* training features. Saves `final_whitelist_fold_N.json`.
-    *   **Scale Features (`scale_features`):** Fits `StandardScaler` on fold's *raw* training features (numeric only). Scales train, val, test features (`X_train_scaled`, etc.). Saves `feature_scaler_fold_N.joblib`.
-    *   **Prune Features (`select_and_prune_features` - Pruning Part):** Prunes the *scaled* data splits (`X_train_scaled` -> `X_train_pruned`) using the `final_whitelist` determined earlier.
-    *   **Create Sequences (`create_sequences`):** Converts the fold's *pruned, scaled* train/val/test sets into sequences (`X_train_seq`, etc.).
-    *   **Train/Load GRU (`train_or_load_gru`):**
-        *   If `sweep_enabled`, runs `GRUHyperTuner` (with updated objective, restricted search space, Keras callback for pruning on `val_loss`, logging objective components) to find best hyperparameters using fold's train/val sequences. Trains final model with best params. Saves best params JSON and Optuna plots.
-        *   If not sweeping, trains/loads GRU using config defaults.
-        *   Saves the final fold GRU model (`gru_model_fold_N.keras`), history, and learning curve plot.
-    *   **Calibrate Probabilities (`calibrate_probabilities`):**
-        *   Fits the selected calibrator (Temp/Vector with optional L2 reg) on the fold's validation sequences.
-        *   If `optimize_edge_threshold`, calculates optimal threshold using Youden's J, stores it internally (`self.optimized_edge_threshold`), and saves it (`optimized_edge_threshold_fold_N.txt`).
-        *   Saves fold calibration parameters (`calibration_{temp/vector}_fold_N.npy`).
-    *   **GRU Validation Gate (`_perform_gru_validation_checks`):** Checks edge-filtered accuracy CI and Brier score against updated config thresholds (`validation_gates.gru`) using the fold's determined `optimized_edge_threshold`. Exits fold if failed.
-    *   **Train/Load SAC (`train_or_load_sac`):**
-        *   If `train_sac`, initializes `SACTrainer` using a config copy (passing the fold's `optimized_edge_threshold` and disabling rolling calibration if active). Trains agent handling PER (with clipping, alpha annealing), Oracle Seeding (with IS weight decay), State Normalization, Action Penalty (`0.01/cost`). Saves agent, filter state, logs, plots in a fold-specific `sac_train_...` dir.
-        *   If loading, determines path from config.
-    *   **Run Backtest (`run_backtest`):**
-        *   Initializes `Backtester`.
-        *   Passes the fold's SAC agent path, test sequences, GRU handler, *initial* calibration state, the fold's `optimized_edge_threshold`, and original test prices.
-        *   If `rolling_enabled`, the backtester uses raw predictions to refit calibrator, potentially triggered early by **ECE Coverage Alarm** (`ECE > config threshold`).
-        *   Logs backtest performance (Sharpe, MDD, Win Rate) and gate pass/fail status. Fold failure here does *not* halt the pipeline immediately but is recorded.
-    *   **Store Fold Results:** Appends the `backtest_metrics` dict (including status) to `all_fold_metrics`. Stores the `sac_agent_load_path` if SAC was trained successfully.
-5.  **Aggregate Metrics (`aggregate_fold_metrics`):** Calculates summary statistics (mean, std, min, max, median) across metrics from all *successful* folds. Saves `aggregated_wf_metrics.json`.
-6.  **Aggregate SAC Agents (`aggregate_sac_agents`):** (Optional: if `sac_aggregation.enabled`)
-    *   Loads SAC agents from the stored paths of successful folds.
-    *   Averages the weights of the loaded agents.
-    *   Saves the aggregated agent to the main run's model dir (`models/run_.../sac_agent_aggregated/`). Saves `sac_aggregation_info.txt`.
-7.  **Final Release Decision (`final_release_decision`):** Evaluates aggregated metrics against overall release criteria defined in `validation_gates.final_release` (e.g., min % successful folds, median Sharpe).
-8.  **Log Final Status:** Logs whether the pipeline passed or failed the final release criteria.
+1.  **Initialization (`run.py`):** Sets up infrastructure (config, logging, IO). Instantiates `TradingPipeline`.
+2.  **Data Loading (`TradingPipeline` calls `data_processing.load_and_preprocess`):** Loads the *entire* raw dataset.
+3.  **Fold Generation (`TradingPipeline._generate_walk_forward_folds`):** Yields date ranges for each fold.
+4.  **Fold Loop (`TradingPipeline.execute`):** Iterates through folds.
+    *   **Select Fold Data:** Extracts raw data for the current fold range.
+    *   **Feature Engineering (`TradingPipeline` calls `data_processing.engineer_features_for_fold`):** Computes features on fold data.
+    *   **Labeling (`TradingPipeline` calls `data_processing.define_labels_and_align_fold`):** Calculates labels.
+    *   **Split Fold Data (`TradingPipeline` calls `data_processing.split_data_fold`):** Splits into `train`, `val`, `test`.
+    *   **Baseline Gate (`TradingPipeline` calls `evaluation.run_baseline_checks_fold`):** Runs Logistic Regression check. Halts fold on failure.
+    *   **Select Features (`TradingPipeline` calls `feature_processing.select_features_fold`):** Performs selection. Saves whitelist.
+    *   **Scale Features (`TradingPipeline` calls `feature_processing.scale_features_fold`):** Fits/applies scaler. Saves scaler.
+    *   **Prune Features (`TradingPipeline` calls `feature_processing.prune_features_fold`):** Prunes scaled data using whitelist.
+    *   **Create Sequences (`TradingPipeline` calls `sequence_creation.create_sequences_fold`):** Creates GRU input sequences.
+    *   **Train/Load GRU (`TradingPipeline` calls `modelling.train_or_load_gru_fold`):** Handles training, Optuna sweep, or loading. Handles re-processing (scale/prune/sequence) internally if loaded scaler differs. Saves model/params.
+    *   **Calibrate Probabilities (`TradingPipeline` calls `modelling.calibrate_probabilities_fold`):** Fits calibrator, optimizes edge threshold. Saves parameters.
+    *   **GRU Validation Gate (`TradingPipeline` calls `evaluation.run_gru_validation_checks_fold`):** Checks calibrated validation predictions. Halts fold on failure.
+    *   **Train/Load SAC (`TradingPipeline` calls `modelling.train_or_load_sac_fold`):** Handles SAC training (calling `SACTrainer`) or determines load path.
+    *   **Run Backtest (`TradingPipeline` calls `evaluation.run_backtest_fold`):** Instantiates `Backtester`, runs simulation, handles rolling calibration, performs backtest validation checks. Halts fold on failure.
+    *   **Store Fold Results:** Appends metrics and SAC agent path (if trained) for aggregation.
+5.  **Aggregate Metrics (`TradingPipeline.aggregate_fold_metrics`):** Calculates summary statistics across successful folds.
+6.  **Aggregate SAC Agents (`TradingPipeline` calls `modelling.aggregate_sac_agents`):** (Optional) Averages weights of successful fold agents.
+7.  **Final Release Decision (`TradingPipeline.final_release_decision`):** Evaluates aggregated metrics against final criteria.
+8.  **Log Final Status:** Reports overall pipeline success/failure.
 
-*(Note: If Walk-Forward is disabled (`walk_forward.enabled=false`), the pipeline runs steps D-T once using static splits based on `split_ratios` config).*
+*(Note: If Walk-Forward is disabled, the pipeline runs the fold processing steps once using static splits).*\
+
+
+## Project Structure
+
+```
+gru_sac_predictor/
+├── config/
+│   └── config.yaml             # Main configuration file
+├── data/                       # Data storage (e.g., parquet files)
+├── logs/                       # Log output directory (run-specific subdirs)
+├── models/                     # Saved models/scalers etc. (run-specific subdirs)
+├── results/                    # Output results, metrics, plots (run-specific subdirs)
+├── src/
+│   ├── pipeline_stages/        # **Refactored stage-specific logic**
+│   │   ├── data_processing.py
+│   │   ├── evaluation.py
+│   │   ├── feature_processing.py
+│   │   ├── modelling.py
+│   │   └── sequence_creation.py
+│   ├── __init__.py
+│   ├── backtester.py           # Backtesting simulation engine
+│   ├── baseline_checker.py     # Baseline logistic regression check
+│   ├── calibrator.py           # Temperature scaling
+│   ├── calibrator_vector.py    # Vector scaling
+│   ├── data_loader.py          # Loads raw data
+│   ├── feature_engineer.py     # Feature creation logic
+│   ├── features.py             # Feature lists/definitions (optional)
+│   ├── gru_hyper_tuner.py      # Optuna hyperparameter tuner for GRU
+│   ├── gru_model_handler.py    # GRU model building, training, loading (v2/v3)
+│   ├── io_manager.py           # Handles file I/O structure
+│   ├── logger_setup.py         # Configures logging
+│   ├── metrics.py              # Performance metrics calculation
+│   ├── sac_agent.py            # SAC agent network definitions
+│   ├── sac_trainer.py          # Offline SAC training orchestration
+│   ├── trading_env.py          # Gym-like environment for SAC training
+│   ├── trading_pipeline.py     # Main pipeline orchestrator class
+│   └── utils/                  # Utility functions (e.g., run_id generation)
+├── tests/                      # Unit/integration tests (optional)
+├── run.py                      # Main execution script
+├── requirements.txt            # Python dependencies
+└── README.md                   # This file
+```
 
 ## Core Components Architecture
 
-This section details the architecture and purpose of key modules.
+*   **`run.py`**: Entry point, sets up IO/logging, runs the pipeline.
+*   **`TradingPipeline` (`src/trading_pipeline.py`):** Orchestrates the workflow by calling stage functions, manages overall state, handles walk-forward loop and validation gates.
+*   **`src/pipeline_stages/*.py`**: Contain the core implementation logic for each distinct step of the pipeline (data processing, feature processing, sequencing, modelling, evaluation).
+*   **`IOManager`, `LoggerSetup`**: Utilities for managing outputs and logging.
+*   **`DataLoader`, `FeatureEngineer`**: Data loading and feature generation.
+*   **`GRUModelHandler`, `GRUHyperTuner`**: GRU model implementation (v2/v3), training, loading, and Optuna tuning.
+*   **`Calibrator`, `VectorCalibrator`**: Probability calibration logic.
+*   **`SACTradingAgent`, `TradingEnv`, `SACTrainer`**: SAC agent definition, training environment, and offline training orchestration (including PER, normalization, seeding).
+*   **`Backtester`, `BaselineChecker`, `metrics.py`**: Backtesting simulation, baseline checks, and performance metric calculations.
 
-### 1. Orchestration & Utilities (`TradingPipeline`, `IOManager`, `LoggerSetup`)
-*   **`TradingPipeline` (`src/trading_pipeline.py`):** The main class coordinating the workflow (data prep, feature eng, model training, calibration, backtesting, aggregation) potentially within a walk-forward loop. Reads config, manages data flow, implements validation gates, and calls other components.
-*   **`IOManager` (`src/io_manager.py`):** Handles standardized file I/O (saving/loading models, dataframes, scalers, configs, plots, reports) within a run-specific directory structure.
-*   **`LoggerSetup` (`src/logger_setup.py`):** Configures Python's `logging` for console and file output with standardized formats and levels.
-
-### 2. Data Handling (`DataLoader`, `FeatureEngineer`)
-*   **`DataLoader` (`src/data_loader.py`):** Loads raw OHLCV data from specified database files/sources.
-*   **`FeatureEngineer` (`src/feature_engineer.py`):** Generates features:
-    *   **TA:** Returns, ATR, EMA, RSI (MACD removed).
-    *   **Cyclical:** Hour, Week (sin/cos).
-    *   **Imbalance:** Chaikin AD, SVI, Gap Imbalance.
-    *   **Micro-structure:** Spread Proxy, Vol-Norm Volume Spike, Return Asymmetry, Close-Location Value, Keltner Band Position (with NaN guards).
-    *   **Leakage Guard:** Uses `shift(1)` on inputs for time-dependent features.
-    *   **Target Definition:** Calculates forward returns and binary/ternary labels.
-    *   **Selection/Pruning:** Performs **leakage check** (`corr(ret+h, feat_t-1)`), then L1+VIF selection on *raw/engineered* data (applying minimal whitelist before VIF), then prunes *scaled* data based on the selection.
-
-### 3. GRU Predictor (`gru_model_handler.py`, `gru_hyper_tuner.py`)
-*   **`gru_model_handler.py`:** **Consolidated GRU implementation.**
-    *   Contains builders for both v2 (`build_gru_model`) and v3 (`build_gru_model_v3`) architectures.
-    *   **v3 Architecture:** GRU -> LayerNorm -> Optional MultiHeadAttention -> GlobalAvgPool -> Dense Heads (`mu`, `dir3_logits`). Includes L2 regularization. Uses Huber loss for `mu`, Focal loss for `dir3`.
-    *   Manages training (with early stopping, CSV logging), saving (.keras format), loading (handles custom loss `gaussian_nll`), and prediction (`predict`, `predict_logits`).
-    *   Selects v2/v3 based on `control.use_v3` flag.
-*   **`gru_hyper_tuner.py`:** Implements Optuna hyperparameter sweep for GRU.
-    *   Called by `TradingPipeline` if sweep is enabled.
-    *   Uses **restricted search space** based on config (`hyperparameter_tuning.gru`).
-    *   Uses combined objective based on config (`objective_metric`, `objective_edge_acc_weight`, `objective_brier_weight`). Logs components to trial attributes.
-    *   Supports pruning via **Keras callback** reporting `val_loss` each epoch.
-    *   Trains final fold model using best found parameters.
-    *   Saves best parameters JSON and Optuna plots per fold.
-
-### 4. Probability Calibration (`Calibrator`, `VectorCalibrator`, `metrics.py`)
-*   **`Calibrator` (`src/calibrator.py`):** Implements Temperature Scaling (learns scalar `T`) with optional L2 regularization.
-*   **`VectorCalibrator` (`src/calibrator_vector.py`):** Implements Vector Scaling (learns matrix `W`, bias `b`) with optional L2 regularization. Preferred for ternary.
-*   **Integration:** `TradingPipeline.calibrate_probabilities` fits chosen calibrator per fold (with L2 reg). If `optimize_edge_threshold=true`, calculates and stores optimal edge for the fold. `Backtester` applies calibration step-by-step, handles rolling recalibration if enabled (with **ECE-based coverage alarm**), using the **static fold calibration** if SAC training was active for the fold.
-*   **Edge Threshold Optimization:** `metrics._calculate_optimal_edge_threshold` finds best threshold using Youden's J. Called by `TradingPipeline` if enabled.
-
-### 5. SAC Agent (`sac_agent.py`, `sac_trainer.py`, `trading_env.py`)
-*   **`sac_agent.py`:** Defines the SAC agent networks (Actor, Critic) and update logic.
-    *   Actor outputs squashed Gaussian distribution parameters.
-    *   Uses twin Q-critics.
-    *   Handles automatic entropy tuning (alpha).
-    *   `train` method accepts a batch and returns losses + TD errors (for PER).
-*   **`trading_env.py`:** Gym-style environment using GRU predictions.
-    *   **State:** `[mu, sigma, edge, |mu|/sigma, position]`.
-    *   Takes action (-1 to +1).
-    *   Calculates reward based on PnL, potentially scaled (`reward_scale`) and penalized (**action penalty**, default: `0.01 / transaction_cost`).
-*   **`sac_trainer.py`:** Orchestrates offline SAC training.
-    *   Loads GRU dependencies for a specific run.
-    *   Prepares validation data for the `TradingEnv` (uses fold's static calibration).
-    *   Initializes `TradingEnv` and `SACTradingAgent`.
-    *   Manages the Replay Buffer:
-        *   Implements `PrioritizedReplayBuffer` if `sac.use_per` is true (with **TD error clipping** and **alpha annealing**). Logs TD error distribution stats.
-        *   Uses `collections.deque` for standard uniform replay.
-        *   Handles **Oracle Seeding** into the buffer (with **IS weight decay** for seeded samples).
-    *   Runs the training loop: interacts with env, stores transitions, samples batches (uniform or PER), updates agent, updates PER priorities.
-    *   Handles **State Normalization** using `MeanStdFilter` if `sac.use_state_filter` is true (saves/loads filter state).
-    *   Saves agent checkpoints, final agent, state filter, and logs (rewards, TensorBoard).
-
-### 6. Evaluation (`Backtester`, `BaselineChecker`, `metrics.py`)
-*   **`Backtester` (`src/backtester.py`):** Evaluates the full system on test data.
-    *   Takes trained GRU, SAC agent, initial calibration state, and the fold's edge threshold.
-    *   Simulates step-by-step trading.
-    *   Applies **rolling calibration** logic if enabled (with ECE check).
-    *   Calculates PnL, equity curve, standard performance metrics.
-    *   Saves detailed results dataframe, metrics summary, and plots per fold.
-    *   Logs performance and whether fold backtest gates passed/failed, but does **not** halt the fold on failure (decision made during final aggregation).
-*   **`BaselineChecker` (`src/baseline_checker.py`):** Performs initial logistic regression check on **raw/engineered** training features.
-*   **`metrics.py`:** Contains calculation functions for Sharpe, Brier, Edge-Filtered Accuracy, ECE, and the Youden's J optimization helper.
 
 ## Configuration (`config.yaml`)
 
-The `config.yaml` file centrally controls the pipeline's behavior. Key sections include:
+The `config.yaml` file centrally controls the pipeline's behavior. See comments within the default `config.yaml` for detailed explanations of each parameter. Key sections include:
 
-*   `base_dirs`: Output directories.
-*   `output`: Figure DPI, size, logging level.
-*   `data`: Data source details, label smoothing.
-*   `features`: Minimal feature whitelist, leakage threshold.
-*   `walk_forward`: Settings for WF validation (enable, days, step). **Note:** `walk_forward.enabled=true` overrides `split_ratios`.
-*   `split_ratios`: Used only if `walk_forward.enabled=false`.
-*   `gru`: General GRU settings (horizon, lookback, ternary flag, flat sigma mult).
-*   `gru_v3`: Specific hyperparameters for the v3 architecture (units, attention, losses, reg).
-*   `hyperparameter_tuning`: Controls Optuna sweep for GRU (enable, trials, timeout, pruning, objective metric/weights).
-*   `calibration`: Method (temp/vector), L2 lambda, optimize edge threshold flag, rolling calibration settings (enable, freq, window, ECE alarm).
-*   `validation_gates`: Thresholds for baseline, GRU gates, and final release decision (median Sharpe, % success).
-*   `sac`: SAC hyperparameters (gamma, tau, LR, alpha, PER settings, oracle seeding, IS weight decay steps, state filter).
-*   `sac_aggregation`: Controls post-run agent averaging (enable, method).
-*   `environment`: Trading env parameters (capital, costs, reward scale, action penalty lambda).
-*   `control`: Flags to enable/disable major stages (train GRU, train SAC, run backtest, use v3, plots), model loading/resuming IDs.
+*   `base_dirs`, `output`: Directory and output settings.
+*   `data`, `features`: Data sources, labeling, feature selection controls.
+*   `walk_forward`, `split_ratios`: Controls walk-forward vs. static splits.
+*   `gru`, `gru_v3`: GRU architecture, training parameters.
+*   `hyperparameter_tuning`: Optuna sweep settings for GRU.
+*   `calibration`: Calibration method, parameters, rolling calibration, ECE alarm.
+*   `validation_gates`: Thresholds for baseline, GRU, backtest, and final release checks.
+*   `sac`, `environment`: SAC agent hyperparameters, PER, seeding, environment settings.
+*   `sac_aggregation`: Agent averaging settings.
+*   `control`: High-level flags (train/load models, enable plots, use v3 GRU).
+
+## Installation
+
+1.  Clone the repository.
+2.  Ensure you have Python 3.8+ installed.
+3.  Set up a virtual environment (recommended):
+    ```bash
+    python -m venv .venv
+    source .venv/bin/activate  # On Windows use `.venv\\Scripts\\activate`
+    ```
+4.  Install dependencies:
+    ```bash
+    pip install -r requirements.txt
+    ```
+    *Note: This installs necessary libraries like TensorFlow, PyTorch, Optuna, scikit-learn, pandas, etc.*
+5.  Prepare your data according to the expected format and update paths in `config.yaml`.
 
 ## Usage
 
-1.  **Setup:** Install requirements (`pip install -r requirements.txt`), prepare data in the specified format/location.
-2.  **Configure:** Edit `config.yaml` (data paths, feature lists, model params, control flags, walk-forward settings, calibration, validation thresholds, tuning, aggregation).
-3.  **Run Pipeline:**
-        ```bash
-    # From project root (develop/gru_sac_predictor/)
-    python gru_sac_predictor/run.py --config path/to/your_config.yaml
+1.  **Configure:** Edit `config.yaml` to set data paths, feature lists, model parameters, control flags (e.g., `train_gru`, `train_sac`), walk-forward settings, validation thresholds, etc.
+2.  **Run Pipeline:** Execute `run.py` from the project root directory (`gru_sac_predictor/`), specifying the configuration file:
+    ```bash
+    # Example execution from the parent directory 'develop/gru_sac_predictor/'
+    python gru_sac_predictor/run.py --config gru_sac_predictor/config/config.yaml
     ```
-4.  **Outputs:** Check `logs/`, `models/`, `results/` directories for run-specific outputs, including fold-specific artifacts if WF is enabled.
+    *   You can use other command-line arguments like `--use-ternary` if implemented in `run.py`.
+3.  **Outputs:** Check the directories specified in `config.yaml` (typically subdirectories within `logs/`, `models/`, `results/`) for run-specific outputs. If Walk-Forward is enabled, you will find fold-specific subdirectories containing models, scalers, plots, and results for each fold.
 
 ## Output Artifacts (Walk-Forward Enabled Example)
 
 *   **Main Run Dirs:** `logs/run_<id>/`, `models/run_<id>/`, `results/run_<id>/`
     *   `run_config.yaml`, `pipeline_<id>.log`
-    *   (Post-run) `aggregated_wf_metrics.json`, `sac_aggregation_info.txt`
-    *   (Post-run) `models/.../sac_agent_aggregated/`
+    *   (Post-run) `aggregated_wf_metrics.json`, `sac_aggregation_info.txt` (if enabled)
+    *   (Post-run) `models/.../sac_agent_aggregated/` (if enabled)
 *   **Fold Dirs (within main run dirs):** e.g., `models/run_<id>/fold_1/`
     *   `models/run_<id>/fold_N/models/`: `gru_model_fold_N.keras`, `calibration_{...}_fold_N.npy`, `feature_scaler_fold_N.joblib`, `final_whitelist_fold_N.json`
     *   `models/run_<id>/fold_N/hypertuning/`: (If sweep enabled) `best_gru_params.json`, Optuna plots.
@@ -248,4 +259,17 @@ The `config.yaml` file centrally controls the pipeline's behavior. Key sections
 
 ## Dependencies
 
-See `requirements.txt`. Key libraries include: TensorFlow, NumPy, Pandas, PyYAML, Scikit-learn, Statsmodels, TA-Lib (via `ta` wrapper), Matplotlib, Seaborn, Optuna, PyTorch (for SAC aggregation). Note `tensorflow-addons` is required for optimal focal loss / attention layers.
+All major Python dependencies are listed in `requirements.txt`. Key libraries include:
+
+*   TensorFlow (for GRU)
+*   PyTorch (for SAC)
+*   Optuna (for hyperparameter tuning)
+*   scikit-learn (for scaling, metrics, baseline)
+*   pandas, numpy
+*   pyyaml
+*   matplotlib, seaborn
+
+Install them using:
+```bash
+pip install -r requirements.txt
+```

gru_sac_predictor/src/pipeline_stages/data_processing.py

@@ -0,0 +1,711 @@
+# Stage functions for loading, initial preprocessing, feature engineering, label generation, and splitting 
+import logging
+import sys # Added for sys.exit
+from datetime import datetime, timezone # Added for datetime
+import pandas as pd
+import numpy as np
+from typing import Tuple, Optional, Any, List, Dict # Added List and Dict
+import matplotlib.pyplot as plt # Added for plotting
+import seaborn as sns # Added for plotting
+
+# --- Component Imports --- #
+# Assuming DataLoader is in the parent directory's src
+# This might need adjustment based on actual project structure
+# Using relative import assuming pipeline_stages is sibling to other src modules
+from ..data_loader import DataLoader, fill_missing_bars
+from ..feature_engineer import FeatureEngineer # Added FeatureEngineer import
+from ..io_manager import IOManager # Added IOManager import
+from ..metrics import calculate_sharpe_ratio # For potential baseline comparison
+
+# --- Local Imports --- #
+# Import the label generation function we moved here
+# Removed duplicate import: from .data_processing import generate_direction_labels
+
+# Assuming tensorflow is installed and available
+try:
+    from tensorflow.keras.utils import to_categorical
+except ImportError:
+    logging.warning("TensorFlow/Keras not found. Ternary label one-hot encoding will fail.")
+    # Define a placeholder if keras is not available
+    def to_categorical(*args, **kwargs):
+        raise NotImplementedError("Keras 'to_categorical' is unavailable.")
+
+logger = logging.getLogger(__name__) # Use module-level logger
+
+# --- Refactored Label Generation Logic (Moved from trading_pipeline.py) --- #
+def generate_direction_labels(df: pd.DataFrame, config: dict) -> Tuple[pd.DataFrame, str, pd.Series, Optional[pd.Series]]:
+    """
+    Calculates forward returns and generates binary, soft binary, or ternary direction labels.
+    Also returns the raw forward returns and the epsilon series used for ternary flat definition.
+
+    Args:
+        df (pd.DataFrame): DataFrame containing at least a 'close' column and DatetimeIndex.
+        config (dict): Pipeline configuration dictionary, expecting keys under 'gru' and 'data'.
+
+    Returns:
+        tuple[pd.DataFrame, str, pd.Series, Optional[pd.Series]]:
+            - DataFrame with added forward return and direction label columns (and NaNs dropped based on labels).
+            - Name of the generated direction label column.
+            - Series containing the calculated forward log returns (`fwd_log_ret`).
+            - Series containing the calculated epsilon (`eps`) threshold if ternary, else None.
+    """
+    if 'close' not in df.columns:
+        raise ValueError("'close' column missing in input DataFrame for label generation.")
+
+    gru_cfg = config.get('gru', {})
+    data_cfg = config.get('data', {})
+    horizon = gru_cfg.get('prediction_horizon', 5)
+    use_ternary = gru_cfg.get('use_ternary', False) # Check if ternary flag is set
+
+    target_ret_col = f'fwd_log_ret_{horizon}'
+    eps_series: Optional[pd.Series] = None # Initialize eps
+
+    # --- Calculate Forward Log Return --- #
+    shifted_close = df['close'].shift(-horizon)
+    fwd_returns = np.log(shifted_close / df['close'])
+    df[target_ret_col] = fwd_returns
+
+    # --- Generate Direction Label (Binary/Soft or Ternary) --- #
+    if use_ternary:
+        k = gru_cfg.get('flat_sigma_multiplier', 0.25)
+        target_dir_col = f'direction_label3_{horizon}'
+        logger.info(f"Generating ternary labels ({target_dir_col}) with k={k}...")
+
+        sigma_n = fwd_returns.rolling(window=horizon, min_periods=max(1, horizon//2)).std()
+        eps = k * sigma_n
+        eps_series = eps # Store the calculated eps series
+
+        conditions = [fwd_returns > eps, fwd_returns < -eps]
+        choices = [2, 0] # 2=up, 0=down
+        ordinal_labels = np.select(conditions, choices, default=1).astype(int) # 1=flat
+
+        # --- Log Distribution & Check Balance --- #
+        df['_ordinal_label_temp'] = ordinal_labels
+        valid_mask_for_dist = ~np.isnan(eps) & ~np.isnan(fwd_returns)
+        ordinal_labels_valid = df.loc[valid_mask_for_dist, '_ordinal_label_temp']
+
+        if not ordinal_labels_valid.empty:
+            counts = np.bincount(ordinal_labels_valid, minlength=3)
+            total_valid = len(ordinal_labels_valid)
+            if total_valid > 0: # Avoid division by zero
+                dist_pct = counts / total_valid * 100
+                log_msg = (f"Label dist (n={total_valid}): "
+                           f"Down(0)={dist_pct[0]:.1f}%, Flat(1)={dist_pct[1]:.1f}%, Up(2)={dist_pct[2]:.1f}%")
+                logger.info(log_msg)
+
+                min_pct_threshold = 10.0 # As per implementation
+                if any(p < min_pct_threshold for p in dist_pct):
+                    error_msg = f"Label imbalance detected! Min class percentage is {np.min(dist_pct):.1f}% (Threshold: {min_pct_threshold}%). Check data or flat_sigma_multiplier (k={k})."
+                    logger.error(error_msg)
+                    print(f"ERROR: {error_msg}") # Also print for visibility
+            else:
+                logger.warning("Label distribution check skipped: total valid labels is zero.")
+        else:
+            logger.warning("Could not calculate label distribution (no valid sigma or returns).")
+        # --- End Distribution Check --- #
+
+        # --- One-hot encode --- #
+        try:
+            y_cat_full = np.full((len(df), 3), np.nan, dtype=np.float32)
+            if not ordinal_labels_valid.empty:
+                 y_cat_valid = to_categorical(ordinal_labels_valid, num_classes=3)
+                 y_cat_full[valid_mask_for_dist] = y_cat_valid.astype(np.float32)
+            else:
+                 logger.warning("No valid ordinal labels to one-hot encode.")
+
+            # Assign the list of arrays (or NaNs) - using list avoids mixed type issues later
+            df[target_dir_col] = [list(row) if not np.all(np.isnan(row)) else np.nan for row in y_cat_full]
+
+        except NotImplementedError as nie:
+             logger.error(f"Ternary label generation failed: {nie}. Keras 'to_categorical' is unavailable. Please install tensorflow.", exc_info=True)
+             raise # Re-raise exception to halt pipeline
+        except Exception as e:
+            logger.error(f"Error during one-hot encoding: {e}", exc_info=True)
+            raise # Re-raise exception to halt pipeline if encoding fails
+        finally:
+            if '_ordinal_label_temp' in df.columns:
+                 df.drop(columns=['_ordinal_label_temp'], inplace=True)
+        # --- End One-hot Encoding --- #
+
+    else: # Binary / Soft Binary
+        target_dir_col = f'direction_label_{horizon}'
+        label_smoothing = data_cfg.get('label_smoothing', 0.0)
+        if not (0.0 <= label_smoothing < 1.0):
+            logger.warning(f"Invalid label_smoothing value ({label_smoothing}). Must be in [0.0, 1.0). Disabling smoothing.")
+            label_smoothing = 0.0
+
+        if label_smoothing > 0.0:
+            high_label = 1.0 - label_smoothing / 2.0
+            low_label = label_smoothing / 2.0
+            logger.info(f"Applying label smoothing: {label_smoothing:.2f} -> labels [{low_label:.2f}, {high_label:.2f}] for {target_dir_col}")
+            df[target_dir_col] = np.where(fwd_returns > 0, high_label, low_label).astype(np.float32)
+        else:
+            logger.info(f"Using hard binary labels (0.0 / 1.0) for {target_dir_col}")
+            df[target_dir_col] = (fwd_returns > 0).astype(np.float32)
+
+    # --- Drop Rows with NaN Targets --- #
+    initial_rows = len(df)
+
+    # Create mask for NaNs in the direction column
+    if use_ternary:
+        # Check if elements are np.nan (since we assign np.nan for rows with no valid labels)
+        nan_mask_dir = df[target_dir_col].isna()
+    else:
+        nan_mask_dir = df[target_dir_col].isna()
+
+    nan_mask_combined = df[target_ret_col].isna() | nan_mask_dir
+
+    df_clean = df[~nan_mask_combined].copy()
+
+    final_rows = len(df_clean)
+    if final_rows < initial_rows:
+        logger.info(f"Dropped {initial_rows - final_rows} rows due to NaN targets (horizon={horizon}).")
+
+    if df_clean.empty:
+         logger.error("DataFrame is empty after defining labels and dropping NaNs. Exiting.")
+         # Returning empty DataFrame, caller should handle exit
+         return pd.DataFrame(), target_dir_col, pd.Series(dtype=float), None # Return empty series/None on failure
+
+    # Return the cleaned df, target col name, and the *original* full fwd_returns and eps series
+    # Need to return the original series aligned with the original df index *before* cleaning
+    # So the caller can align them with the features *after* cleaning df_clean
+    return df_clean, target_dir_col, fwd_returns, eps_series
+# --- End Label Generation --- # 
+
+# --- Stage 1: Load and Preprocess Data (Moved from TradingPipeline.load_and_preprocess_data) --- #
+def load_and_preprocess(
+    data_loader: DataLoader,
+    io: Optional[IOManager],
+    run_id: str,
+    config: Dict[str, Any]
+) -> Tuple[Optional[pd.DataFrame], Optional[Dict[str, Any]]]:
+    """
+    Loads the full raw dataset using DataLoader and performs initial checks.
+
+    Args:
+        data_loader: Initialized DataLoader instance.
+        io: IOManager instance (optional).
+        run_id: Current run ID.
+        config: Pipeline configuration dictionary.
+
+    Returns:
+        Tuple containing:
+        - DataFrame with raw loaded data, or None on failure.
+        - Dictionary summarizing the loading process, or None on failure.
+    """
+    logger.info("--- Stage: Loading and Preprocessing Data ---")
+    data_cfg = config.get('data', {})
+    
+    # --- Extract necessary parameters from config --- #
+    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
+    vol_sampling = data_cfg.get('volatility_sampling', {}).get('enabled', False)
+    vol_window = data_cfg.get('volatility_sampling', {}).get('window', 30)
+    vol_quantile = data_cfg.get('volatility_sampling', {}).get('quantile', 0.5)
+    
+    # Validate required parameters
+    if not all([ticker, exchange, start_date, end_date]):
+        logger.error("Missing required data parameters in config: ticker, exchange, start_date, end_date")
+        return None, None
+    # --- End Parameter Extraction --- #
+
+    load_summary = {
+        'ticker': ticker,
+        'exchange': exchange,
+        'start_date_req': start_date,
+        'end_date_req': end_date,
+        'interval_req': interval,
+        'vol_sampling_enabled': vol_sampling,
+        'vol_window': vol_window,
+        'vol_quantile': vol_quantile,
+    }
+
+    try:
+        logger.info(f"Loading data for {ticker} ({exchange}) from {start_date} to {end_date}, interval {interval}")
+        # --- Pass extracted parameters to load_data --- #
+        df_raw = data_loader.load_data(
+            ticker=ticker,
+            exchange=exchange,
+            start_date=start_date,
+            end_date=end_date,
+            interval=interval,
+            vol_sampling=vol_sampling,
+            vol_window=vol_window,
+            vol_quantile=vol_quantile
+        )
+        # --- End Pass Parameters --- #
+
+        if df_raw is None or df_raw.empty:
+            logger.error("Data loading returned empty DataFrame or failed.")
+            return None, load_summary
+
+        # --- Fill Missing Bars (Step 2.5 from prompts/missing_data.txt) --- #
+        if io is None:
+            logger.error("IOManager is required for fill_missing_bars reporting. Cannot proceed.")
+            return None, load_summary
+        try:
+            df_filled = fill_missing_bars(df_raw, config, io, logger)
+            if df_filled is None or df_filled.empty:
+                logger.error("fill_missing_bars returned empty DataFrame or failed.")
+                return None, load_summary
+            df_raw = df_filled # Replace df_raw with the filled version
+            logger.info("Missing bars handled successfully.")
+        except ValueError as e:
+             logger.error(f"Error during missing bar handling: {e}. Halting processing.")
+             return None, load_summary
+        except Exception as e:
+             logger.error(f"Unexpected error during missing bar handling: {e}. Halting processing.", exc_info=True)
+             return None, load_summary
+        # --- End Fill Missing Bars --- #
+
+        # Calculate memory usage and log info
+        mem_usage = df_raw.memory_usage(deep=True).sum() / (1024**2)
+        if load_summary:
+            logger.info(f"Data loading summary: {load_summary}")
+        else:
+            logger.warning("No load summary returned by DataLoader.")
+        logger.info(f"Loaded data: {df_raw.shape[0]} rows, {df_raw.shape[1]} columns. Memory: {mem_usage:.2f} MB")
+        logger.info(f"Time range: {df_raw.index.min()} to {df_raw.index.max()}")
+
+        # --- V3 Output Contract: Stage 1 Artifacts --- #
+        if io:
+            if load_summary:
+                save_summary = load_summary.copy() # Don't modify original
+                save_summary['run_id'] = run_id
+                save_summary['timestamp_utc'] = datetime.now(timezone.utc).isoformat()
+                # TODO: Finalize summary content (add counts, NaN info etc.)
+                logger.info("Saving preprocess summary...")
+                io.save_json(save_summary, "preprocess_summary", use_txt=True) # Spec wants .txt
+            
+            # Save head of preprocessed data
+            if df_raw is not None and not df_raw.empty:
+                 logger.info("Saving head of preprocessed data (first 20 rows)...")
+                 io.save_df(df_raw.head(20), "head_preprocessed") 
+            else:
+                 logger.warning("Skipping saving head_preprocessed: DataFrame is empty or None.")
+
+        else:
+            logger.warning("IOManager not available, skipping saving of Stage 1 artifacts (preprocess_summary, head_preprocessed).")
+        # --- End V3 Output Contract ---
+
+        # --- V3 Output Contract: Stage 2 Artifact (Label Histogram) --- #
+        # TODO: Move this plotting logic to evaluation stage or after split, needs y_train.
+        # if io and config.get('control', {}).get('generate_plots', True):
+        #     logger.info("Generating training label distribution histogram... [SKIPPED IN CURRENT STAGE]")
+            # ... (Original plotting code removed from here)
+        # --- End V3 Output Contract ---
+
+        return df_raw, load_summary 
+
+    except Exception as e:
+        logger.error(f"Error during data loading: {e}", exc_info=True)
+        return None, None
+
+# --- Stage 2: Engineer Features (Moved from TradingPipeline.engineer_features) --- #
+def engineer_features_for_fold(
+    df: pd.DataFrame,
+    feature_engineer: FeatureEngineer,
+    io: Optional[IOManager], # Added IOManager for saving figure
+    config: Dict[str, Any], # Added config for plot settings
+    target_col: Optional[str] = None # Added target column name for sorting correlation
+) -> pd.DataFrame:
+    """Adds features using FeatureEngineer, handles NaNs, and saves correlation heatmap for a fold.
+
+    Args:
+        df (pd.DataFrame): Input DataFrame for the fold (typically raw data).
+        feature_engineer (FeatureEngineer): Initialized FeatureEngineer instance.
+        io (Optional[IOManager]): IOManager instance for saving artifacts.
+        config (Dict[str, Any]): Pipeline configuration dictionary.
+        target_col (Optional[str]): Name of the target column to sort correlations by (e.g., 'fwd_log_ret_5').
+
+    Returns:
+        pd.DataFrame: DataFrame with engineered features, NaNs dropped.
+                     Returns an empty DataFrame if input is empty or result is empty.
+    """
+    logger.info("--- Stage: Engineering Features --- ")
+    if df is None or df.empty:
+        logger.error("Input DataFrame is empty. Cannot engineer features.")
+        return pd.DataFrame() # Return empty DataFrame to indicate failure
+
+    if feature_engineer is None:
+        logger.error("FeatureEngineer not initialized. Cannot engineer features.")
+        # Or raise an error? For now return empty
+        return pd.DataFrame()
+
+    # Add base features (cyclical, imbalance, TA)
+    df_engineered = feature_engineer.add_base_features(df.copy())
+
+    # --- V3 Output Contract: Feature Correlation Heatmap --- #
+    # Generate heatmap *before* dropping NaNs to capture full feature set correlations
+    # if io and config.get('control', {}).get('generate_plots', True): # Check if plotting is enabled
+    if io: # Assume generate_plots is implicitly true if io is provided
+        try:
+            logger.info("Generating feature correlation heatmap...")
+            numeric_cols = df_engineered.select_dtypes(include=np.number).columns
+            if len(numeric_cols) < 2:
+                 logger.warning("Skipping correlation heatmap: Less than 2 numeric columns found.")
+            else:
+                corr_matrix = df_engineered[numeric_cols].corr(method='pearson')
+
+                # Get plot settings from config
+                output_cfg = config.get('output', {})
+                fig_size = output_cfg.get('figure_size', [16, 9])
+                plot_style = output_cfg.get('plot_style', 'seaborn-v0_8-darkgrid')
+                annot_threshold = output_cfg.get('corr_annot_threshold', 0.5)
+                plot_footer = output_cfg.get('plot_footer', "© GRU-SAC v3")
+
+                plt.style.use(plot_style)
+                fig, ax = plt.subplots(figsize=fig_size)
+
+                sort_features = False
+                if target_col and target_col in corr_matrix.columns:
+                    # Sort by absolute correlation with the target
+                    target_corr = corr_matrix[target_col].abs().sort_values(ascending=False)
+                    sorted_cols = target_corr.index.tolist()
+                    corr_matrix_sorted = corr_matrix.loc[sorted_cols, sorted_cols]
+                    sort_features = True
+                else:
+                    if target_col:
+                         logger.warning(f"Target column '{target_col}' not found in correlation matrix. Heatmap will not be sorted by target correlation.")
+                    corr_matrix_sorted = corr_matrix # Use original matrix if no target or not found
+
+                sns.heatmap(
+                    corr_matrix_sorted,
+                    annot=False, # Annotations can be messy; spec only requires > threshold
+                    cmap='coolwarm', # Diverging palette centered at 0
+                    center=0,
+                    linewidths=0.5,
+                    cbar=True,
+                    square=True, # Ensure square cells
+                    ax=ax
+                )
+
+                # Annotate cells where absolute correlation > threshold (from config)
+                for i in range(corr_matrix_sorted.shape[0]):
+                    for j in range(corr_matrix_sorted.shape[1]):
+                        if abs(corr_matrix_sorted.iloc[i, j]) > annot_threshold and i != j:
+                            ax.text(j + 0.5, i + 0.5, f'{corr_matrix_sorted.iloc[i, j]:.2f}',
+                                    ha='center', va='center', color='black', fontsize=8)
+
+                title = "Feature Correlation Heatmap (Pearson)"
+                if sort_features:
+                    title += f" - Sorted by |ρ| vs '{target_col}'"
+                ax.set_title(title, fontsize=14)
+                plt.xticks(rotation=90, fontsize=8)
+                plt.yticks(rotation=0, fontsize=8)
+
+                # Add footer (from config)
+                if plot_footer: # Only add if footer is not empty
+                    plt.figtext(0.99, 0.01, plot_footer, ha="right", va="bottom", fontsize=8, color='gray')
+
+                # Save figure using IOManager
+                io.save_figure(fig, "feature_corr_heatmap", section='figures') # Saved to results/<run_id>/figures/
+                plt.close(fig) # Close figure after saving
+                logger.info("Saved feature correlation heatmap.")
+
+        except Exception as e:
+            logger.error(f"Failed to generate or save feature correlation heatmap: {e}", exc_info=True)
+    else:
+        logger.warning("IOManager not provided or plotting disabled, skipping feature correlation heatmap.")
+    # --- End V3 Output Contract --- #
+
+    # --- REMOVE Aggressive DropNA --- #
+    # Dropping all rows with any NaN here is too aggressive, especially with long lookback features.
+    # NaN handling should occur within feature calculation methods (bfill/ffill/fillna(0))
+    # and critically during label definition (where rows without valid labels are dropped).
+    # initial_rows = len(df_engineered)
+    # df_engineered.dropna(inplace=True)
+    # rows_dropped = initial_rows - len(df_engineered)
+    # if rows_dropped > 0:
+    #      logger.warning(f"Dropped {rows_dropped} rows with NaN values after feature engineering.")
+    # --- End REMOVE --- #
+
+    # Check if dataframe became empty *after feature calculation and internal NaN handling*
+    # (Though ideally internal handling should prevent this)
+    if df_engineered.empty:
+         logger.error("DataFrame is empty after feature engineering (check internal NaN handling in FeatureEngineer)." )
+         return pd.DataFrame() # Return empty DataFrame
+
+    logger.info(f"Feature engineering complete. Shape: {df_engineered.shape}")
+    return df_engineered 
+
+# --- Stage 3: Define Labels and Align (Moved from TradingPipeline.define_labels_and_align) --- #
+def define_labels_and_align_fold(
+    df_engineered: pd.DataFrame, 
+    config: dict
+) -> Tuple[pd.DataFrame, str, List[str], pd.Series, Optional[pd.Series]]:
+    """Defines prediction labels, aligns with features, and separates targets for a fold.
+    Also returns the raw forward returns and epsilon series used for filtering baselines.
+
+    Args:
+        df_engineered (pd.DataFrame): DataFrame with engineered features for the fold.
+        config (dict): Pipeline configuration dictionary.
+
+    Returns:
+        Tuple[pd.DataFrame, str, List[str], pd.Series, Optional[pd.Series]]:
+            - df_labeled_aligned: DataFrame with labels generated and features/targets aligned (NaNs dropped).
+            - target_dir_col: Name of the direction label column.
+            - target_cols: List containing names of all target columns (ret + dir).
+            - fwd_returns_aligned: Series of forward returns aligned with df_labeled_aligned.
+            - eps_aligned: Series of epsilon threshold aligned with df_labeled_aligned (or None).
+            Returns (pd.DataFrame(), "", [], pd.Series(), None) on failure or empty input.
+    """
+    logger.info("--- Stage: Defining Labels and Aligning --- ")
+    if df_engineered is None or df_engineered.empty:
+        logger.error("Engineered data (DataFrame) is empty. Cannot define labels.")
+        return pd.DataFrame(), "", [], pd.Series(dtype=float), None
+
+    # --- Call the label generation function (already in this module) --- #
+    try:
+        # generate_direction_labels modifies the DataFrame in place and returns it
+        # It also returns the original fwd_returns and eps series (aligned with df_engineered)
+        df_clean, target_dir_col, fwd_returns_orig, eps_orig = generate_direction_labels(
+            df_engineered.copy(), # Pass a copy to avoid modifying original outside this scope if needed
+            config
+        )
+    except Exception as e:
+         logger.error(f"Label generation failed: {e}.", exc_info=True)
+         return pd.DataFrame(), "", [], pd.Series(dtype=float), None
+
+    if df_clean.empty:
+        logger.error("Label generation resulted in an empty DataFrame.")
+        return pd.DataFrame(), "", [], pd.Series(dtype=float), None
+    # --- End Label Generation Call --- #
+
+    # --- Determine Target Columns --- #
+    horizon = config.get('gru', {}).get('prediction_horizon', 5)
+    target_ret_col = f'fwd_log_ret_{horizon}'
+    # target_dir_col is returned by generate_direction_labels
+    target_cols = [target_ret_col, target_dir_col]
+
+    # Ensure the columns actually exist after generation and cleaning
+    if not all(col in df_clean.columns for col in target_cols):
+         # Log which columns are actually present for debugging
+         present_cols = df_clean.columns.tolist()
+         logger.error(f"Generated label/return columns ({target_cols}) not found in DataFrame after label generation. Present columns: {present_cols}")
+         return pd.DataFrame(), "", [], pd.Series(dtype=float), None
+    # --- End Determine Target Columns --- #
+
+    # --- Align fwd_returns_orig and eps_orig with the cleaned DataFrame --- #
+    fwd_returns_aligned = fwd_returns_orig.loc[df_clean.index]
+    eps_aligned = eps_orig.loc[df_clean.index] if eps_orig is not None else None
+    # --- End Alignment --- #
+
+    # Note: Separation of X and y happens in the splitting function now
+    # We just need to return the fully labeled/aligned DataFrame and target column names.
+    logger.info(f"Labels defined and aligned. Shape: {df_clean.shape}")
+
+    # Return the aligned DataFrame and the aligned supplementary series
+    return df_clean, target_dir_col, target_cols, fwd_returns_aligned, eps_aligned 
+
+# --- Stage 4: Split Data (Moved from TradingPipeline.split_data) --- #
+def split_data_fold(
+    df_labeled_aligned: pd.DataFrame,
+    fwd_returns_aligned: pd.Series,
+    eps_aligned: Optional[pd.Series],
+    config: dict,
+    target_columns: List[str],
+    target_dir_col: str,
+    fold_dates: Optional[Tuple] = None,
+    current_fold: Optional[int] = None # For logging
+) -> Tuple[
+    # Features
+    pd.DataFrame, pd.DataFrame, pd.DataFrame,
+    # Original Targets
+    pd.DataFrame, pd.DataFrame, pd.DataFrame,
+    # Original Full DataFrames
+    pd.DataFrame, pd.DataFrame, pd.DataFrame,
+    # Ordinal Direction Target (Train only)
+    pd.Series,
+    # Forward Returns (Train/Val)
+    pd.Series, Optional[pd.Series],
+    # Epsilon (Train/Val)
+    Optional[pd.Series], Optional[pd.Series],
+    # Ordinal Direction Labels (Val)
+    Optional[pd.Series]
+]:
+    """Splits features, targets, fwd returns, and epsilon for a given fold.
+
+    Args:
+        df_labeled_aligned (pd.DataFrame): Labeled and aligned data for the entire fold period.
+        fwd_returns_aligned (pd.Series): Forward returns aligned with df_labeled_aligned.
+        eps_aligned (Optional[pd.Series]): Epsilon threshold aligned with df_labeled_aligned.
+        config (dict): Pipeline configuration.
+        target_columns (List[str]): Names of all target columns (e.g., ['fwd_log_ret_5', 'direction_label_5']).
+        target_dir_col (str): Name of the specific direction target column.
+        fold_dates (Optional[Tuple]): Tuple of (train_start, train_end, val_start, val_end, test_start, test_end) for WF.
+        current_fold (Optional[int]): Fold number for logging.
+
+    Returns:
+        Tuple containing the split dataframes/series:
+        (X_train_raw, X_val_raw, X_test_raw, # Features
+         y_train, y_val, y_test, # Targets
+         df_train_original, df_val_original, df_test_original, # Original DFs
+         y_dir_train_ordinal, # Ordinal Direction Labels (Train)
+         fwd_ret_train, fwd_ret_val, # Forward Returns (Train/Val)
+         eps_train, eps_val, # Epsilon (Train/Val, Optional)
+         y_dir_val_ordinal) # Ordinal Direction Labels (Val, Optional)
+         Returns tuple of Nones if splitting fails.
+    """
+    fold_label = f"Fold {current_fold}" if current_fold is not None else "Split"
+    logger.info(f"--- {fold_label}: Stage: Splitting Data --- ")
+
+    if df_labeled_aligned is None or df_labeled_aligned.empty:
+        logger.error(f"Fold {fold_label}: Input data for splitting is empty.")
+        # Return Nones to indicate failure (update count based on new returns)
+        return (None,) * 14
+
+    # --- Temporarily add fwd_ret and eps to DataFrame for easier splitting --- #
+    temp_fwd_ret_col = '__temp_fwd_ret__'
+    temp_eps_col = '__temp_eps__'
+    df_split_input = df_labeled_aligned.copy()
+    df_split_input[temp_fwd_ret_col] = fwd_returns_aligned
+    if eps_aligned is not None:
+        df_split_input[temp_eps_col] = eps_aligned
+    # --- End Temp Add --- #
+
+    if not isinstance(df_split_input.index, pd.DatetimeIndex):
+         logger.error(f"{fold_label}: Data index must be DatetimeIndex for splitting. Aborting.")
+         raise SystemExit(f"{fold_label}: Index is not DatetimeIndex in split_data.")
+
+    if not target_columns:
+         logger.error(f"{fold_label}: Target columns list is empty. Aborting.")
+         raise SystemExit(f"{fold_label}: Target columns missing in split_data.")
+    if not target_dir_col:
+         logger.error(f"{fold_label}: Target direction column name is empty. Aborting.")
+         raise SystemExit(f"{fold_label}: Target direction column missing in split_data.")
+
+    # Ensure target columns exist before trying to drop/select them
+    cols_to_drop = [col for col in target_columns if col in df_split_input.columns]
+    if len(cols_to_drop) != len(target_columns):
+        missing_targets = set(target_columns) - set(cols_to_drop)
+        logger.error(f"{fold_label}: Expected target columns {missing_targets} not found in input DataFrame. Aborting.")
+        raise SystemExit(f"{fold_label}: Missing target columns in split_data input.")
+
+    # Exclude temporary columns from feature_cols
+    feature_cols = df_split_input.columns.difference(cols_to_drop + [temp_fwd_ret_col, temp_eps_col])
+    if feature_cols.empty:
+        logger.error(f"{fold_label}: No feature columns remain after excluding targets and temp cols. Aborting.")
+        raise SystemExit(f"{fold_label}: No feature columns found in split_data.")
+
+    # --- Determine if ternary mode is active --- #
+    use_ternary = config.get('gru', {}).get('use_ternary', False)
+    # --- End Determine Ternary --- #
+
+    # Initialize split results
+    X_train_raw, X_val_raw, X_test_raw = pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
+    y_train, y_val, y_test = pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
+    df_train_original, df_val_original, df_test_original = pd.DataFrame(), pd.DataFrame(), pd.DataFrame()
+    fwd_ret_train, fwd_ret_val = pd.Series(dtype=float), pd.Series(dtype=float)
+    eps_train, eps_val = None, None
+    y_dir_train_raw_format, y_dir_val_raw_format = pd.Series(dtype=object), pd.Series(dtype=object) # Store raw labels before converting
+
+    # Split based on Walk-Forward dates or ratios
+    if fold_dates and len(fold_dates) == 6 and all(fold_dates): # Check for valid WF tuple
+        train_start, train_end, val_start, val_end, test_start, test_end = fold_dates
+        logger.info(f"  Splitting using Walk-Forward dates: Train=[{train_start}, {train_end}), Val=[{val_start}, {val_end}), Test=[{test_start}, {test_end})")
+
+        # Slicing logic
+        df_train_original = df_split_input.loc[train_start:train_end]
+        df_val_original = df_split_input.loc[val_start:val_end]
+        df_test_original = df_split_input.loc[test_start:test_end] if test_start else pd.DataFrame()
+
+    else: # Single split using ratios
+        split_cfg = config.get('split_ratios', {})
+        train_ratio = split_cfg.get('train', 0.7)
+        val_ratio = split_cfg.get('validation', 0.15)
+        test_ratio = round(1.0 - train_ratio - val_ratio, 2)
+        logger.info(f"  Splitting using ratios: Train={train_ratio:.2f}, Val={val_ratio:.2f}, Test={test_ratio:.2f}")
+
+        total_len = len(df_split_input)
+        train_end_idx = int(total_len * train_ratio)
+        val_end_idx = int(total_len * (train_ratio + val_ratio))
+
+        df_train_original = df_split_input.iloc[:train_end_idx]
+        df_val_original = df_split_input.iloc[train_end_idx:val_end_idx]
+        df_test_original = df_split_input.iloc[val_end_idx:]
+
+    # --- Extract components from split DataFrames --- #
+    if not df_train_original.empty:
+        X_train_raw = df_train_original[feature_cols]
+        y_train = df_train_original[target_columns]
+        y_dir_train_raw_format = df_train_original[target_dir_col]
+        fwd_ret_train = df_train_original[temp_fwd_ret_col]
+        if temp_eps_col in df_train_original:
+            eps_train = df_train_original[temp_eps_col]
+
+    if not df_val_original.empty:
+        X_val_raw = df_val_original[feature_cols]
+        y_val = df_val_original[target_columns]
+        fwd_ret_val = df_val_original[temp_fwd_ret_col]
+        if temp_eps_col in df_val_original:
+            eps_val = df_val_original[temp_eps_col]
+
+    if not df_test_original.empty:
+        X_test_raw = df_test_original[feature_cols]
+        y_test = df_test_original[target_columns]
+    # --- End Extraction --- #
+
+    # --- Extract Ordinal Labels if Ternary --- #
+    y_dir_train_ordinal = None
+    if not y_train.empty: # Check if training data exists
+        if use_ternary:
+            valid_mask = y_dir_train_raw_format.notna() & y_dir_train_raw_format.apply(lambda x: isinstance(x, list) and len(x) == 3)
+            if valid_mask.any():
+                 ordinal_values = y_dir_train_raw_format[valid_mask].apply(np.argmax)
+                 y_dir_train_ordinal = pd.Series(np.nan, index=y_dir_train_raw_format.index)
+                 y_dir_train_ordinal[valid_mask] = ordinal_values
+                 logger.info(f"{fold_label}: Extracted ordinal labels (0, 1, 2) for feature selection. Count: {valid_mask.sum()}")
+            else:
+                 logger.warning(f"{fold_label}: No valid list-based ternary labels found in y_dir_train_raw_format to convert to ordinal.")
+                 y_dir_train_ordinal = pd.Series(dtype=np.float64) # Return empty series
+        else:
+            y_dir_train_ordinal = y_dir_train_raw_format.astype(int) # Ensure integer type
+    else:
+        y_dir_train_ordinal = pd.Series(dtype=int) # Empty series if no train data
+    # --- End Extract Ordinal Labels --- #
+
+    # --- Extract Ordinal Validation Labels if Ternary --- #
+    y_dir_val_ordinal = None
+    if not y_val.empty: # Check if validation data exists
+        if use_ternary:
+            # Use y_dir_val_raw_format which holds the lists/NaNs
+            valid_mask_val = y_dir_val_raw_format.notna() & y_dir_val_raw_format.apply(lambda x: isinstance(x, list) and len(x) == 3)
+            if valid_mask_val.any():
+                 ordinal_values_val = y_dir_val_raw_format[valid_mask_val].apply(np.argmax)
+                 y_dir_val_ordinal = pd.Series(np.nan, index=y_dir_val_raw_format.index)
+                 y_dir_val_ordinal[valid_mask_val] = ordinal_values_val
+                 logger.info(f"{fold_label}: Extracted ordinal validation labels. Count: {valid_mask_val.sum()}")
+            else:
+                 logger.warning(f"{fold_label}: No valid ternary labels found in y_dir_val_raw_format.")
+                 y_dir_val_ordinal = pd.Series(dtype=np.float64)
+        else:
+            # Use y_dir_val_raw_format which holds 0.0/1.0
+            y_dir_val_ordinal = y_dir_val_raw_format.astype(int)
+    else:
+        y_dir_val_ordinal = pd.Series(dtype=int) # Empty series if no validation data
+    # --- End Extract Ordinal Validation Labels --- #
+
+    # Log split shapes and check for empty splits
+    logger.info(f"Data split complete for {fold_label}:")
+    logger.info(f"  Train: X={X_train_raw.shape}, y={y_train.shape}, fwd_ret={fwd_ret_train.shape}, eps={eps_train.shape if eps_train is not None else 'None'} ({X_train_raw.index.min()} to {X_train_raw.index.max()})" if not X_train_raw.empty else "  Train: EMPTY")
+    logger.info(f"  Val:   X={X_val_raw.shape}, y={y_val.shape}, fwd_ret={fwd_ret_val.shape}, eps={eps_val.shape if eps_val is not None else 'None'} ({X_val_raw.index.min()} to {X_val_raw.index.max()})" if not X_val_raw.empty else "  Val:   EMPTY")
+    logger.info(f"  Test:  X=({X_test_raw.shape if X_test_raw is not None else 'None'}), y=({y_test.shape if y_test is not None else 'None'}) ({df_test_original.index.min() if df_test_original is not None and not df_test_original.empty else 'N/A'} to {df_test_original.index.max() if df_test_original is not None and not df_test_original.empty else 'N/A'})" )
+
+    # Check required splits are non-empty
+    if X_train_raw.empty or X_val_raw.empty:
+        logger.error(f"Fold {current_fold}: Data splitting resulted in empty train or validation set. Aborting fold.")
+        raise SystemExit(f"Fold {current_fold}: Empty train or validation split detected.")
+
+    return (
+        X_train_raw, X_val_raw, X_test_raw, # Features
+        y_train, y_val, y_test, # Targets
+        df_train_original, df_val_original, df_test_original, # Original DFs
+        y_dir_train_ordinal, # Ordinal Direction Labels (Train)
+        fwd_ret_train, fwd_ret_val, # Forward Returns (Train/Val)
+        eps_train, eps_val, # Epsilon (Train/Val, Optional)
+        y_dir_val_ordinal # Ordinal Direction Labels (Val, Optional)
+    )

gru_sac_predictor/src/pipeline_stages/sequence_creation.py

@@ -0,0 +1,182 @@
+# Stage functions for creating GRU input sequences
+import logging
+import sys
+import numpy as np
+import pandas as pd
+from typing import Tuple, Dict, Optional, List
+import json # Added for saving artefact
+
+# Assuming IOManager is importable from parent src directory
+from ..io_manager import IOManager 
+
+logger = logging.getLogger(__name__)
+
+def create_sequences_fold(
+    X_data: pd.DataFrame,
+    y_data: pd.DataFrame,
+    target_names: List[str], # e.g., ['mu', 'dir3'] or ['mu', 'dir']
+    lookback: int,
+    name: str, # e.g., "Train", "Validation", "Test"
+    config: dict, # For gru.drop_imputed_sequences
+    io: Optional[IOManager] # For saving artefact
+) -> Tuple[Optional[np.ndarray], Optional[Dict], Optional[pd.Index], int]:
+    """
+    Transforms pruned, scaled feature DataFrame into 3D sequences for GRU input
+    and extracts corresponding targets for a specific data split (Train/Val/Test).
+    Handles dropping sequences containing imputed bars based on config.
+
+    Args:
+        X_data (pd.DataFrame): Pruned, scaled features for the split.
+        y_data (pd.DataFrame): Targets for the split.
+        target_names (List[str]): List of target column names in y_data (e.g., ['mu', 'dir3']).
+        lookback (int): Sequence length.
+        name (str): Name of the split (e.g., "Train", "Validation", "Test") for logging.
+        config (dict): Pipeline configuration dictionary.
+        io (Optional[IOManager]): IOManager instance for saving artefacts.
+
+    Returns:
+        Tuple containing:
+        - X_seq (np.ndarray or None): 3D feature sequences.
+        - y_seq_dict (Dict or None): Dictionary of target sequences.
+        - target_indices (pd.Index or None): Timestamps corresponding to the targets.
+        - dropped_count (int): Number of sequences dropped due to imputed bars.
+        Returns (None, None, None, 0) if sequence creation fails or data is insufficient.
+        Raises SystemExit on critical errors (e.g., misalignment).
+    """
+    logger.info(f"--- Creating {name} Sequences ---")
+    use_ternary = config.get('gru', {}).get('use_ternary', False)
+    drop_imputed = config.get('gru', {}).get('drop_imputed_sequences', False)
+    imputed_col_name = 'bar_imputed' # Assuming this is the column name
+
+    # --- Input Validation --- #
+    if X_data is None or y_data is None or X_data.empty or y_data.empty:
+        logger.error(f"{name}: Missing or empty features/targets for sequence creation.")
+        return None, None, None, 0
+
+    # Check for bar_imputed column
+    if imputed_col_name not in X_data.columns:
+         logger.error(f"{name}: Required column '{imputed_col_name}' not found in features. Cannot handle imputed sequences.")
+         # Decide whether to proceed without it or raise error - raising for now
+         raise SystemExit(f"{name}: '{imputed_col_name}' column missing. Sequence creation halted.")
+
+    # Strict Anti-Leakage Check
+    try:
+        assert X_data.index.equals(y_data.index), \
+            f"{name}: Features and targets indices misaligned!"
+    except AssertionError as e:
+        logger.error(f"Data alignment check failed: {e}. Potential data leakage. Aborting.")
+        raise SystemExit(f"{name}: {e}")
+
+    # Check target columns exist
+    if not all(col in y_data.columns for col in target_names):
+         missing_targets = set(target_names) - set(y_data.columns)
+         logger.error(f"{name}: Target columns {missing_targets} not found in y_data. Aborting.")
+         raise SystemExit(f"{name}: Missing target columns for sequencing.")
+    # --- End Input Validation --- #
+
+    # Convert DataFrames to numpy for potential speedup, keep index access
+    features_np = X_data.values
+    imputed_flag_np = X_data[imputed_col_name].values.astype(bool) # Ensure boolean type
+    # Extract targets based on target_names
+    targets_dict_np = {name: y_data[name].values for name in target_names}
+
+    X_seq_list, y_seq_dict_list = [], {name: [] for name in target_names}
+    mask_seq_list = [] # To store the imputed flag sequences
+    target_indices = []
+
+    if len(X_data) <= lookback:
+         logger.warning(f"{name}: DataFrame length ({len(X_data)}) is not greater than lookback ({lookback}). Cannot create sequences.")
+         return None, None, None, 0
+
+    for i in range(lookback, len(features_np)):
+        # Feature window: [i-lookback, i)
+        X_seq_list.append(features_np[i - lookback : i])
+        mask_seq_list.append(imputed_flag_np[i - lookback : i])
+
+        # Targets correspond to index i
+        for t_name in target_names:
+             target_val = targets_dict_np[t_name][i]
+             # Special handling for potential list/array type in ternary labels
+             if use_ternary and 'dir' in t_name and isinstance(target_val, list):
+                 target_val = np.array(target_val, dtype=np.float32)
+             y_seq_dict_list[t_name].append(target_val)
+             
+        target_indices.append(y_data.index[i]) # Get index corresponding to target
+
+    if not X_seq_list: # Check if any sequences were created
+         logger.warning(f"{name}: No sequences were generated (length <= lookback?).")
+         return None, None, None, 0
+
+    # Convert lists to numpy arrays
+    X_seq = np.array(X_seq_list, dtype=np.float32)
+    mask_seq = np.array(mask_seq_list, dtype=bool)
+    target_indices_pd = pd.Index(target_indices)
+    y_seq_dict_np = {}
+    for t_name in target_names:
+        try:
+            # Attempt to stack; requires consistent shapes
+            if use_ternary and 'dir' in t_name:
+                 y_seq_dict_np[t_name] = np.stack(y_seq_dict_list[t_name]).astype(np.float32)
+            else: # Assuming other targets are scalar
+                 y_seq_dict_np[t_name] = np.array(y_seq_dict_list[t_name], dtype=np.float32)
+        except ValueError as e:
+            logger.error(f"{name}: Error stacking target '{t_name}': {e}. Check target consistency (especially ternary).", exc_info=True)
+            shapes = [getattr(item, 'shape', type(item)) for item in y_seq_dict_list[t_name]]
+            from collections import Counter
+            logger.error(f"Target shapes/types found: {Counter(shapes)}")
+            raise SystemExit(f"{name}: Inconsistent target shapes for '{t_name}' during sequence creation.") from e
+
+    orig_n = X_seq.shape[0]
+    dropped_count = 0
+
+    # Conditionally drop sequences containing imputed bars
+    if drop_imputed:
+        logger.info(f"{name}: Dropping sequences containing imputed bars (drop_imputed_sequences=True)...")
+        valid_mask = ~mask_seq.any(axis=1)
+        X_seq = X_seq[valid_mask]
+        mask_seq = mask_seq[valid_mask] # Keep mask aligned, though not explicitly used later
+        for t_name in target_names:
+             y_seq_dict_np[t_name] = y_seq_dict_np[t_name][valid_mask]
+        target_indices_pd = target_indices_pd[valid_mask]
+        
+        dropped_count = orig_n - X_seq.shape[0]
+        logger.info(f"{name}: Generated {orig_n} sequences, dropped {dropped_count} containing imputed bars. Remaining: {X_seq.shape[0]}")
+
+        # Save summary artifact
+        if io:
+             summary_data = {
+                 "split_name": name,
+                 "total_sequences_generated": orig_n,
+                 "sequences_dropped_imputed": dropped_count,
+                 "sequences_remaining": X_seq.shape[0],
+                 "drop_imputed_sequences_config": drop_imputed
+             }
+             try:
+                 filename = f"imputed_sequence_summary_{name.lower()}.json"
+                 io.save_json(summary_data, filename, section='results', indent=4)
+                 logger.info(f"Saved imputed sequence summary to results/{filename}")
+             except Exception as e:
+                 logger.error(f"Failed to save imputed sequence summary for {name}: {e}")
+        else:
+             logger.warning(f"IOManager not available, cannot save imputed sequence summary for {name}.")
+             
+    else:
+         logger.info(f"{name}: Generated {orig_n} sequences. Keeping sequences with imputed bars (drop_imputed_sequences=False).")
+
+    # Final checks
+    if X_seq.shape[0] == 0:
+         logger.error(f"{name}: No valid sequences remaining after potential filtering. Aborting.")
+         return None, None, None, dropped_count # Return 0 count if no sequences left
+
+    # --- REMOVE: Final Dictionary Mapping (Let GRU handler manage this) --- #
+    # final_y_seq_dict = {
+    #     'mu': y_seq_dict_np['ret'], # Map 'ret' to 'mu'
+    #     'dir3': y_seq_dict_np['dir3'] # Keep 'dir3' as is
+    # }
+    # --- END REMOVE --- #
+
+    # Log final shapes
+    logger.info(f"Sequence shapes created for {name}:")
+    logger.info(f"  X={X_seq.shape}, y_keys={list(y_seq_dict_np.keys())}, indices={len(target_indices_pd)}")
+
+    return X_seq, y_seq_dict_np, target_indices_pd, dropped_count

gru_sac_predictor/src/trading_env.py

@@ -6,6 +6,8 @@ Uses pre-calculated GRU predictions (mu, sigma, p_cal) and actual returns.
 import numpy as np
 import pandas as pd
 import logging
+import gymnasium as gym
+from omegaconf import DictConfig # Added for config typing
 
 env_logger = logging.getLogger(__name__)
 
@@ -15,6 +17,8 @@ class TradingEnv:
                  sigma_predictions: np.ndarray,
                  p_cal_predictions: np.ndarray,
                  actual_returns: np.ndarray,
+                 bar_imputed_flags: np.ndarray, # Added imputed flags
+                 config: DictConfig, # Added config
                  initial_capital: float = 10000.0,
                  transaction_cost: float = 0.0005,
                  reward_scale: float = 100.0,
@@ -27,18 +31,22 @@ class TradingEnv:
             sigma_predictions: Predicted volatility (σ̂ = exp(log σ̂)).
             p_cal_predictions: Calibrated probability of price increase (p_cal).
             actual_returns: Actual log returns (y_ret).
+            bar_imputed_flags: Boolean array indicating if a bar was imputed.
+            config: OmegaConf configuration object.
             initial_capital: Starting capital for simulation (used notionally in reward).
             transaction_cost: Fractional cost per trade.
             reward_scale: Multiplier for the reward signal.
             action_penalty_lambda: Coefficient for the action magnitude penalty (λ).
         """
-        assert len(mu_predictions) == len(sigma_predictions) == len(p_cal_predictions) == len(actual_returns), \
-            "All input arrays must have the same length"
+        assert len(mu_predictions) == len(sigma_predictions) == len(p_cal_predictions) == len(actual_returns) == len(bar_imputed_flags), \
+            "All input arrays (predictions, returns, imputed_flags) must have the same length"
 
         self.mu = mu_predictions
         self.sigma = sigma_predictions
         self.p_cal = p_cal_predictions
         self.actual_returns = actual_returns
+        self.bar_imputed = bar_imputed_flags.astype(bool) # Store imputed flags
+        self.config = config # Store config
 
         self.initial_capital = initial_capital
         self.transaction_cost = transaction_cost
@@ -65,20 +73,36 @@ class TradingEnv:
         self.state_dim = 5
         self.action_dim = 1
 
+        # --- Define Gym Spaces ---
+        self.action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(self.action_dim,), dtype=np.float32)
+        self.observation_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.state_dim,), dtype=np.float32)
+        # --- End Define Gym Spaces ---
+
         env_logger.info(f"TradingEnv initialized with {self.n_steps} steps.")
 
     def _get_state(self) -> np.ndarray:
         """Construct the state vector for the current step."""
         if self.current_step >= self.n_steps:
-             # Handle episode end - return a dummy state or zeros
              return np.zeros(self.state_dim, dtype=np.float32)
 
         mu_t = self.mu[self.current_step]
         sigma_t = self.sigma[self.current_step]
         p_cal_t = self.p_cal[self.current_step]
 
-        edge_t = 2 * p_cal_t - 1
-        z_score_t = np.abs(mu_t) / (sigma_t + 1e-9)
+        # Calculate edge based on p_cal shape (binary vs ternary)
+        if isinstance(p_cal_t, (np.ndarray, list)) and len(p_cal_t) == 3:
+            # Ternary: edge = max(P(up), P(down)) - P(flat)
+            # Assuming order [Down, Flat, Up] for p_cal_t
+            edge_t = max(p_cal_t[2], p_cal_t[0]) - p_cal_t[1]
+        elif isinstance(p_cal_t, (float, np.number)):
+            # Binary: edge = 2 * P(up) - 1
+            edge_t = 2 * p_cal_t - 1
+        else:
+            env_logger.error(f"Unexpected type/shape for p_cal_t at step {self.current_step}: {p_cal_t}. Using edge=0.")
+            edge_t = 0.0
+
+        _EPS = 1e-9 # Define epsilon locally
+        z_score_t = np.abs(mu_t) / (sigma_t + _EPS)
 
         # State uses position *before* the action for this step is taken
         state = np.array([
@@ -108,11 +132,48 @@ class TradingEnv:
         Returns:
             tuple: (next_state, reward, done, info_dict)
         """
+        info = {'capital': self.current_capital, 'position': self.current_position, 'is_imputed_step_skipped': False}
+
         if self.current_step >= self.n_steps:
             # Should not happen if 'done' is handled correctly, but as safeguard
             env_logger.warning("Step called after environment finished.")
-            return self._get_state(), 0.0, True, {}
+            return self._get_state(), 0.0, True, info
 
+        # --- Handle Imputed Bar --- #
+        imputed = self.bar_imputed[self.current_step]
+        if imputed:
+            mode = self.config.sac.imputed_handling
+            env_logger.debug(f"SAC step {self.current_step} on imputed bar: handling={mode}")
+            if mode == "skip":
+                self.current_step += 1
+                next_state = self._get_state() # Get state for the *next* actual step
+                # Return 0 reward, not done, but indicate skip for buffer handling
+                info['is_imputed_step_skipped'] = True 
+                return next_state, 0.0, False, info
+            elif mode == "hold":
+                # Action is forced to maintain current position
+                action = self.current_position 
+            elif mode == "penalty":
+                # Calculate reward penalty based on config
+                target_position_penalty = np.clip(action, -1.0, 1.0)
+                reward = -self.config.sac.action_penalty * (target_position_penalty - self.current_position)**2
+                # Update position based on agent's intended action (clipped)
+                self.current_position = target_position_penalty
+                # Update capital notionally (no actual return, only cost if implemented)
+                # Cost is implicitly 0 here as there's no trade size if pos doesn't change
+                # If penalty mode allowed position change, cost would apply.
+                # For simplicity, we don't add cost here for the penalty step.
+                self.current_step += 1
+                next_state = self._get_state()
+                scaled_reward = reward * self.reward_scale # Scale the penalty
+                done = self.current_step >= self.n_steps
+                info['capital'] = self.current_capital
+                info['position'] = self.current_position
+                return next_state, scaled_reward, done, info
+            # else: default behavior (treat as normal bar) - implicitly handled by falling through
+        # --- End Handle Imputed Bar --- #
+
+        # --- Normal Step Logic (if not imputed or handling mode allows fallthrough like 'hold') --- #
         # Action is the TARGET position for the *end* of this step
         target_position = np.clip(action, -1.0, 1.0)
         trade_size = target_position - self.current_position
@@ -150,7 +211,9 @@ class TradingEnv:
         done = self.current_step >= self.n_steps or self.current_capital <= 0
 
         next_state = self._get_state()
-        info = {'capital': self.current_capital, 'position': self.current_position}
+        # Update info dict (capital/position might have changed in normal step)
+        info['capital'] = self.current_capital
+        info['position'] = self.current_position 
 
         # Log step details periodically
         # if self.current_step % 1000 == 0:

gru_sac_predictor/tests/test_sequence_creation.py

@@ -0,0 +1,186 @@
+import pytest
+import pandas as pd
+import numpy as np
+from omegaconf import OmegaConf
+from unittest.mock import MagicMock
+import os
+import tempfile
+import json
+
+# Adjust the import path based on your project structure
+from gru_sac_predictor.src.pipeline_stages.sequence_creation import create_sequences_fold
+from gru_sac_predictor.src.io_manager import IOManager # Adjust path if needed
+
+# --- Test Fixtures ---
+
+@pytest.fixture
+def sample_data_with_imputed():
+    """Creates sample X and y dataframes with a 'bar_imputed' column."""
+    dates = pd.to_datetime(pd.date_range('2023-01-01', periods=20, freq='T'))
+    lookback = 5
+    n_features_orig = 3
+    n_samples = len(dates)
+
+    # Features (including bar_imputed)
+    X_data = pd.DataFrame(
+        np.random.randn(n_samples, n_features_orig), 
+        index=dates, 
+        columns=[f'feat_{i}' for i in range(n_features_orig)]
+    )
+    # Add bar_imputed column - mark some bars as imputed
+    imputed_flags = np.zeros(n_samples, dtype=bool)
+    imputed_flags[2] = True # Imputed within first potential sequence
+    imputed_flags[8] = True # Imputed within a later potential sequence
+    imputed_flags[15] = True # Imputed near the end
+    X_data['bar_imputed'] = imputed_flags
+
+    # Targets (mu and dir3)
+    y_data = pd.DataFrame({
+        'mu': np.random.randn(n_samples),
+        'dir3': [list(row) for row in np.eye(3)[np.random.randint(0, 3, n_samples)]] # Example one-hot
+    }, index=dates)
+
+    return X_data, y_data
+
+@pytest.fixture
+def base_config():
+    """Creates a base OmegaConf config for testing sequence creation."""
+    conf = OmegaConf.create({
+        'gru': {
+            'lookback': 5,
+            'use_ternary': True, # Matches sample_data_with_imputed
+            'drop_imputed_sequences': True # Default to True for testing dropping
+        },
+        # Add other necessary sections if needed
+    })
+    return conf
+
+@pytest.fixture
+def mock_io_manager():
+    """Creates a mock IOManager for testing artefact saving."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        mock_io = MagicMock(spec=IOManager)
+        mock_io.results_dir = tmpdir
+        saved_jsons = {}
+        def mock_save_json(data, filename, **kwargs):
+            filepath = os.path.join(tmpdir, filename)
+            saved_jsons[filename] = data
+            with open(filepath, 'w') as f:
+                json.dump(data, f, **kwargs)
+        mock_io.save_json.side_effect = mock_save_json
+        mock_io.get_artifact_path.side_effect = lambda filename: os.path.join(tmpdir, filename)
+        mock_io._saved_jsons = saved_jsons 
+        yield mock_io
+
+# --- Test Functions ---
+
+def test_sequence_creation_shapes(sample_data_with_imputed, base_config, mock_io_manager):
+    X_data, y_data = sample_data_with_imputed
+    lookback = base_config.gru.lookback
+    n_features = X_data.shape[1]
+    n_samples = len(X_data)
+    expected_n_seq = n_samples - lookback
+
+    # Test without dropping imputed
+    cfg_no_drop = base_config.copy()
+    cfg_no_drop.gru.drop_imputed_sequences = False
+
+    X_seq, y_seq_dict, indices, dropped_count = create_sequences_fold(
+        X_data=X_data, y_data=y_data, target_names=['mu', 'dir3'], 
+        lookback=lookback, name="TestSplit", config=cfg_no_drop, io=mock_io_manager
+    )
+
+    assert dropped_count == 0
+    assert X_seq is not None
+    assert y_seq_dict is not None
+    assert indices is not None
+    assert X_seq.shape == (expected_n_seq, lookback, n_features)
+    assert 'mu' in y_seq_dict and y_seq_dict['mu'].shape == (expected_n_seq,)
+    assert 'dir3' in y_seq_dict and y_seq_dict['dir3'].shape == (expected_n_seq, 3)
+    assert len(indices) == expected_n_seq
+    # Check first target index corresponds to lookback-th original index
+    assert indices[0] == X_data.index[lookback]
+    # Check last target index corresponds to last original index
+    assert indices[-1] == X_data.index[-1]
+
+def test_sequence_dropping_imputed(sample_data_with_imputed, base_config, mock_io_manager):
+    X_data, y_data = sample_data_with_imputed
+    lookback = base_config.gru.lookback
+    n_samples = len(X_data)
+    expected_n_seq_orig = n_samples - lookback
+
+    # Config with dropping enabled (default in fixture)
+    cfg_drop = base_config
+
+    X_seq, y_seq_dict, indices, dropped_count = create_sequences_fold(
+        X_data=X_data.copy(), y_data=y_data.copy(), target_names=['mu', 'dir3'], 
+        lookback=lookback, name="TestDrop", config=cfg_drop, io=mock_io_manager
+    )
+
+    assert X_seq is not None
+    assert y_seq_dict is not None
+    assert indices is not None
+    
+    # Determine which original sequences should have been dropped
+    # A sequence starting at index i uses data from [i, i+lookback-1]
+    # The target corresponds to index i+lookback
+    # We need to check the imputed flag in the range [i, i+lookback-1] for each potential sequence target index i+lookback
+    
+    # Original target indices range from index `lookback` to `n_samples - 1`
+    should_be_dropped_mask = np.zeros(expected_n_seq_orig, dtype=bool)
+    imputed_flags_np = X_data['bar_imputed'].values
+    for seq_idx in range(expected_n_seq_orig):
+        # The features for this sequence are from original indices [seq_idx, seq_idx + lookback - 1]
+        feature_indices_range = slice(seq_idx, seq_idx + lookback)
+        if np.any(imputed_flags_np[feature_indices_range]):
+            should_be_dropped_mask[seq_idx] = True
+            
+    expected_dropped_count = np.sum(should_be_dropped_mask)
+    expected_remaining_count = expected_n_seq_orig - expected_dropped_count
+
+    assert dropped_count == expected_dropped_count
+    assert X_seq.shape[0] == expected_remaining_count
+    assert y_seq_dict['mu'].shape[0] == expected_remaining_count
+    assert y_seq_dict['dir3'].shape[0] == expected_remaining_count
+    assert len(indices) == expected_remaining_count
+    
+    # Check that the remaining indices are correct (weren't marked for dropping)
+    original_indices = X_data.index[lookback:]
+    expected_remaining_indices = original_indices[~should_be_dropped_mask]
+    pd.testing.assert_index_equal(indices, expected_remaining_indices)
+
+    # Check artifact saving
+    assert 'imputed_sequence_summary_testdrop.json' in mock_io_manager._saved_jsons
+    report_data = mock_io_manager._saved_jsons['imputed_sequence_summary_testdrop.json']
+    assert report_data['total_sequences_generated'] == expected_n_seq_orig
+    assert report_data['sequences_dropped_imputed'] == expected_dropped_count
+    assert report_data['sequences_remaining'] == expected_remaining_count
+
+def test_sequence_creation_no_imputed_col(sample_data_with_imputed, base_config, mock_io_manager):
+    X_data, y_data = sample_data_with_imputed
+    X_data_no_imputed = X_data.drop(columns=['bar_imputed'])
+    lookback = base_config.gru.lookback
+
+    with pytest.raises(SystemExit) as excinfo:
+        create_sequences_fold(
+            X_data=X_data_no_imputed, y_data=y_data, target_names=['mu', 'dir3'], 
+            lookback=lookback, name="TestNoImputedCol", config=base_config, io=mock_io_manager
+        )
+    assert "'bar_imputed' column missing" in str(excinfo.value)
+
+def test_sequence_creation_insufficient_data(sample_data_with_imputed, base_config, mock_io_manager):
+    X_data, y_data = sample_data_with_imputed
+    lookback = base_config.gru.lookback
+    # Create data shorter than lookback
+    X_short = X_data.iloc[:lookback-1]
+    y_short = y_data.iloc[:lookback-1]
+
+    X_seq, y_seq_dict, indices, dropped_count = create_sequences_fold(
+        X_data=X_short, y_data=y_short, target_names=['mu', 'dir3'], 
+        lookback=lookback, name="TestShort", config=base_config, io=mock_io_manager
+    )
+
+    assert X_seq is None
+    assert y_seq_dict is None
+    assert indices is None
+    assert dropped_count == 0 

gru_sac_predictor/tests/test_trading_env.py

@@ -0,0 +1,166 @@
+import pytest
+import numpy as np
+from omegaconf import OmegaConf
+
+# Adjust import path based on structure
+from gru_sac_predictor.src.trading_env import TradingEnv
+
+# --- Test Fixtures ---
+
+@pytest.fixture
+def sample_env_data():
+    """Provides sample data for initializing the TradingEnv."""
+    n_steps = 10
+    data = {
+        'mu_predictions': np.random.randn(n_steps) * 0.001,
+        'sigma_predictions': np.abs(np.random.randn(n_steps) * 0.002 + 0.005),
+        'p_cal_predictions': np.random.rand(n_steps),
+        'actual_returns': np.random.randn(n_steps) * 0.0015,
+        'bar_imputed_flags': np.array([False, False, True, False, True, True, False, False, True, False], dtype=bool)
+    }
+    return data
+
+@pytest.fixture
+def base_env_config():
+    """Base configuration for the environment."""
+    return OmegaConf.create({
+        'sac': {
+            'imputed_handling': 'skip', # Default test mode
+            'action_penalty': 0.05
+        },
+        'environment': {
+            'initial_capital': 10000.0,
+            'transaction_cost': 0.0005,
+            'reward_scale': 100.0,
+            'action_penalty_lambda': 0.0 # Usually overridden by transaction_cost calc
+        }
+    })
+
+@pytest.fixture
+def trading_env_instance(sample_env_data, base_env_config):
+    """Creates a TradingEnv instance with default 'skip' mode."""
+    return TradingEnv(**sample_env_data, config=base_env_config)
+
+# --- Test Functions ---
+
+def test_env_initialization(trading_env_instance, sample_env_data):
+    assert trading_env_instance.n_steps == len(sample_env_data['actual_returns'])
+    assert trading_env_instance.current_step == 0
+    assert trading_env_instance.current_position == 0.0
+    assert np.array_equal(trading_env_instance.bar_imputed, sample_env_data['bar_imputed_flags'])
+
+def test_env_reset(trading_env_instance):
+    # Take a few steps
+    trading_env_instance.step(0.5)
+    trading_env_instance.step(-0.2)
+    assert trading_env_instance.current_step > 0
+    # Reset
+    initial_state = trading_env_instance.reset()
+    assert trading_env_instance.current_step == 0
+    assert trading_env_instance.current_position == 0.0
+    assert initial_state is not None
+    assert initial_state.shape == (trading_env_instance.state_dim,)
+
+def test_env_step_normal(trading_env_instance):
+    # Test a normal step (step 0 is not imputed)
+    initial_pos = trading_env_instance.current_position
+    action = 0.7
+    next_state, reward, done, info = trading_env_instance.step(action)
+
+    assert trading_env_instance.current_step == 1
+    assert trading_env_instance.current_position == action # Position updates to action
+    assert not info['is_imputed_step_skipped']
+    assert not done
+    assert next_state is not None
+    # Reward calculation is complex, just check type/sign if needed
+    assert isinstance(reward, float)
+
+def test_env_step_imputed_skip(trading_env_instance, sample_env_data):
+    # Step 2 is imputed in sample_env_data
+    trading_env_instance.step(0.5) # Step 0
+    trading_env_instance.step(0.6) # Step 1
+    assert trading_env_instance.current_step == 2
+    initial_pos_before_imputed = trading_env_instance.current_position
+
+    # Action for the imputed step (should be ignored by 'skip')
+    action_imputed = 0.9
+    next_state, reward, done, info = trading_env_instance.step(action_imputed)
+
+    # Should skip step 2 and now be at step 3
+    assert trading_env_instance.current_step == 3 
+    # Position should NOT have changed from step 1
+    assert trading_env_instance.current_position == initial_pos_before_imputed
+    assert reward == 0.0 # Skip gives 0 reward
+    assert not done
+    assert info['is_imputed_step_skipped'] == True # Crucial check for buffer
+    # Check that the returned state is for step 3
+    expected_state_step_3 = trading_env_instance._get_state() # Get state now that we are at step 3
+    np.testing.assert_array_almost_equal(next_state, expected_state_step_3)
+
+def test_env_step_imputed_hold(sample_env_data, base_env_config):
+    cfg = base_env_config.copy()
+    cfg.sac.imputed_handling = 'hold'
+    env = TradingEnv(**sample_env_data, config=cfg)
+
+    # Step 2 is imputed
+    env.step(0.5) # Step 0
+    env.step(0.6) # Step 1
+    assert env.current_step == 2
+    position_before_imputed = env.current_position
+
+    # Action for the imputed step (should be overridden by 'hold')
+    action_imputed = -0.5 
+    next_state, reward, done, info = env.step(action_imputed)
+
+    # Should process step 2 and move to step 3
+    assert env.current_step == 3
+    # Position should be the same as before the step
+    assert env.current_position == position_before_imputed
+    assert not info['is_imputed_step_skipped']
+    assert not done
+    # Reward should be calculated based on holding the position
+    expected_pnl = position_before_imputed * (np.exp(sample_env_data['actual_returns'][2]) - 1)
+    expected_cost = 0 # No trade size if holding
+    expected_penalty = 0 # No penalty in hold mode
+    expected_raw_reward = expected_pnl - expected_cost - expected_penalty
+    expected_scaled_reward = expected_raw_reward * cfg.environment.reward_scale
+    assert np.isclose(reward, expected_scaled_reward)
+
+def test_env_step_imputed_penalty(sample_env_data, base_env_config):
+    cfg = base_env_config.copy()
+    cfg.sac.imputed_handling = 'penalty'
+    cfg.sac.action_penalty = 0.1 # Use a specific penalty for testing
+    env = TradingEnv(**sample_env_data, config=cfg)
+
+    # Step 2 is imputed
+    env.step(0.5) # Step 0
+    env.step(0.6) # Step 1
+    assert env.current_step == 2
+    position_before_imputed = env.current_position # Should be 0.6
+
+    # Action for the imputed step
+    action_imputed = -0.2 
+    next_state, reward, done, info = env.step(action_imputed)
+
+    # Should process step 2 and move to step 3
+    assert env.current_step == 3
+    # Position should update to the *agent's* action
+    assert env.current_position == np.clip(action_imputed, -1.0, 1.0)
+    assert not info['is_imputed_step_skipped']
+    assert not done
+    
+    # Reward calculation is ONLY the penalty
+    expected_raw_reward = -cfg.sac.action_penalty * (action_imputed - position_before_imputed)**2
+    expected_scaled_reward = expected_raw_reward * cfg.environment.reward_scale
+    assert np.isclose(reward, expected_scaled_reward)
+
+def test_env_done_condition(trading_env_instance, sample_env_data):
+    n_steps = len(sample_env_data['actual_returns'])
+    # Step through the environment
+    done = False
+    for i in range(n_steps):
+        _, _, done, _ = trading_env_instance.step(np.random.uniform(-1, 1))
+        if i < n_steps - 1:
+            assert not done
+        else:
+            assert done # Should be done on the last step