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`.
+*   **Modular Pipeline Structure:** Pipeline logic is refactored into stage-specific functions within `src/pipeline_stages/` for improved readability, maintainability, and testability (see Project Structure).
-*   **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.
+*   **Consolidated GRU Logic:** Both v2 and v3 GRU model architectures are implemented and managed within `src/gru_model_handler.py`.
-*   **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`.
+*   **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.
+    *   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`). Uses the **static** calibration from training time if SAC training is active to prevent lookahead.
+    *   **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.
+*   **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** (default: `0.01 / transaction_cost`), oracle seeding with **Importance Sampling (IS) weight decay** (`per_seed_decay_steps`).
+*   **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:
+*   **Refined Validation Gates:** Configurable thresholds (`validation_gates`) for baseline checks, GRU validation, fold backtest performance, and final release decisions.
-    *   **Baseline Gate:** Checks Logistic Regression CI on **raw/engineered** training features before scaling.
+*   **Micro-structure Features:** Added bar-level features (`FeatureEngineer._add_microstructure_features`) with NaN guards.
-    *   **GRU Gate:** Checks Edge Acc CI and Brier score on validation set after calibration, using the fold's **determined edge threshold**.
+*   **Leakage Guard:** Feature calculations use `shift(1)`. Selection includes correlation check against future returns. Minimal whitelist applied before VIF.
-    *   **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.
+*   **Configuration:** Centralized and expanded `config.yaml`.
-*   **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).
 *   **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];
+            F --> G[data_processing: Engineer Features];
-            G --> H[Define Labels & Align];
+            G --> H[data_processing: Define Labels & Align];
-            H --> I[Split Fold Data];
+            H --> I[data_processing: Split Fold Data];
-            I --> J1[Baseline Check];
+            I --> J1[evaluation: Baseline Check];
-            J1 -- Pass --> L[Select Features];
+            J1 -- Pass --> L[feature_processing: Select Features];
-            L --> K[Scale Features];
+            L --> K[feature_processing: Scale Features];
-            K --> M[Prune Scaled Features];
+            K --> M[feature_processing: Prune Scaled Features];
-            M --> N[Create Sequences];
+            M --> N[sequence_creation: Create Sequences];
-            N --> O[Train/Load GRU];
+            N --> O[modelling: Train/Load GRU];
-            O --> P[Calibrate Probabilities];
+            O --> P[modelling: Calibrate Probabilities];
-            P --> R[GRU Validation Gate];
+            P --> R[evaluation: GRU Validation Gate];
-            R -- Pass --> S[Train/Load SAC Agent];
+            R -- Pass --> S[modelling: Train/Load SAC Agent];
-            S --> T[Run Backtest];
+            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`):**
+1.  **Initialization (`run.py`):** Sets up infrastructure (config, logging, IO). Instantiates `TradingPipeline`.
-    *   Parses args (`--config`, etc.).
+2.  **Data Loading (`TradingPipeline` calls `data_processing.load_and_preprocess`):** Loads the *entire* raw dataset.
-    *   Loads `config.yaml`.
+3.  **Fold Generation (`TradingPipeline._generate_walk_forward_folds`):** Yields date ranges for each fold.
-    *   Initializes `IOManager`, `LoggerSetup`.
+4.  **Fold Loop (`TradingPipeline.execute`):** Iterates through folds.
-    *   Instantiates `TradingPipeline`.
+    *   **Select Fold Data:** Extracts raw data for the current fold range.
-2.  **Data Loading (`TradingPipeline.load_and_preprocess_data`):**
+    *   **Feature Engineering (`TradingPipeline` calls `data_processing.engineer_features_for_fold`):** Computes features on fold data.
-    *   Loads the *entire* raw dataset specified in the config.
+    *   **Labeling (`TradingPipeline` calls `data_processing.define_labels_and_align_fold`):** Calculates labels.
-3.  **Fold Generation (`TradingPipeline._generate_walk_forward_folds`):**
+    *   **Split Fold Data (`TradingPipeline` calls `data_processing.split_data_fold`):** Splits into `train`, `val`, `test`.
-    *   Based on `walk_forward` config (train/val/test/step days), yields date ranges for each fold.
+    *   **Baseline Gate (`TradingPipeline` calls `evaluation.run_baseline_checks_fold`):** Runs Logistic Regression check. Halts fold on failure.
-4.  **Fold Loop (`TradingPipeline.execute`):** Iterates through generated folds.
+    *   **Select Features (`TradingPipeline` calls `feature_processing.select_features_fold`):** Performs selection. Saves whitelist.
-    *   **Select Fold Data:** Extracts raw data corresponding to the current fold's (Train+Val+Test) date range.
+    *   **Scale Features (`TradingPipeline` calls `feature_processing.scale_features_fold`):** Fits/applies scaler. Saves scaler.
-    *   **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).
+    *   **Prune Features (`TradingPipeline` calls `feature_processing.prune_features_fold`):** Prunes scaled data using whitelist.
-    *   **Labeling (`define_labels_and_align`):** Calculates forward returns and target labels (binary/ternary) for the fold's engineered data.
+    *   **Create Sequences (`TradingPipeline` calls `sequence_creation.create_sequences_fold`):** Creates GRU input sequences.
-    *   **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`.
+    *   **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.
-    *   **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`.
+    *   **Calibrate Probabilities (`TradingPipeline` calls `modelling.calibrate_probabilities_fold`):** Fits calibrator, optimizes edge threshold. Saves parameters.
-    *   **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`.
+    *   **GRU Validation Gate (`TradingPipeline` calls `evaluation.run_gru_validation_checks_fold`):** Checks calibrated validation predictions. Halts fold on failure.
-    *   **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`.
+    *   **Train/Load SAC (`TradingPipeline` calls `modelling.train_or_load_sac_fold`):** Handles SAC training (calling `SACTrainer`) or determines load path.
-    *   **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.
+    *   **Run Backtest (`TradingPipeline` calls `evaluation.run_backtest_fold`):** Instantiates `Backtester`, runs simulation, handles rolling calibration, performs backtest validation checks. Halts fold on failure.
-    *   **Create Sequences (`create_sequences`):** Converts the fold's *pruned, scaled* train/val/test sets into sequences (`X_train_seq`, etc.).
+    *   **Store Fold Results:** Appends metrics and SAC agent path (if trained) for aggregation.
-    *   **Train/Load GRU (`train_or_load_gru`):**
+5.  **Aggregate Metrics (`TradingPipeline.aggregate_fold_metrics`):** Calculates summary statistics across successful folds.
-        *   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.
+6.  **Aggregate SAC Agents (`TradingPipeline` calls `modelling.aggregate_sac_agents`):** (Optional) Averages weights of successful fold agents.
-        *   If not sweeping, trains/loads GRU using config defaults.
+7.  **Final Release Decision (`TradingPipeline.final_release_decision`):** Evaluates aggregated metrics against final criteria.
-        *   Saves the final fold GRU model (`gru_model_fold_N.keras`), history, and learning curve plot.
+8.  **Log Final Status:** Reports overall pipeline success/failure.
-    *   **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.
 
-*(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.
+*   `base_dirs`, `output`: Directory and output settings.
-*   `output`: Figure DPI, size, logging level.
+*   `data`, `features`: Data sources, labeling, feature selection controls.
-*   `data`: Data source details, label smoothing.
+*   `walk_forward`, `split_ratios`: Controls walk-forward vs. static splits.
-*   `features`: Minimal feature whitelist, leakage threshold.
+*   `gru`, `gru_v3`: GRU architecture, training parameters.
-*   `walk_forward`: Settings for WF validation (enable, days, step). **Note:** `walk_forward.enabled=true` overrides `split_ratios`.
+*   `hyperparameter_tuning`: Optuna sweep settings for GRU.
-*   `split_ratios`: Used only if `walk_forward.enabled=false`.
+*   `calibration`: Calibration method, parameters, rolling calibration, ECE alarm.
-*   `gru`: General GRU settings (horizon, lookback, ternary flag, flat sigma mult).
+*   `validation_gates`: Thresholds for baseline, GRU, backtest, and final release checks.
-*   `gru_v3`: Specific hyperparameters for the v3 architecture (units, attention, losses, reg).
+*   `sac`, `environment`: SAC agent hyperparameters, PER, seeding, environment settings.
-*   `hyperparameter_tuning`: Controls Optuna sweep for GRU (enable, trials, timeout, pruning, objective metric/weights).
+*   `sac_aggregation`: Agent averaging settings.
-*   `calibration`: Method (temp/vector), L2 lambda, optimize edge threshold flag, rolling calibration settings (enable, freq, window, ECE alarm).
+*   `control`: High-level flags (train/load models, enable plots, use v3 GRU).
-*   `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).
+## Installation
-*   `sac_aggregation`: Controls post-run agent averaging (enable, method).
+
-*   `environment`: Trading env parameters (capital, costs, reward scale, action penalty lambda).
+1.  Clone the repository.
-*   `control`: Flags to enable/disable major stages (train GRU, train SAC, run backtest, use v3, plots), model loading/resuming IDs.
+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.
+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.  **Configure:** Edit `config.yaml` (data paths, feature lists, model params, control flags, walk-forward settings, calibration, validation thresholds, tuning, aggregation).
+2.  **Run Pipeline:** Execute `run.py` from the project root directory (`gru_sac_predictor/`), specifying the configuration file:
-3.  **Run Pipeline:**
     ```bash
-    # From project root (develop/gru_sac_predictor/)
+    # Example execution from the parent directory 'develop/gru_sac_predictor/'
-    python gru_sac_predictor/run.py --config path/to/your_config.yaml
+    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) `aggregated_wf_metrics.json`, `sac_aggregation_info.txt` (if enabled)
-    *   (Post-run) `models/.../sac_agent_aggregated/`
+    *   (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/sac_trainer.py

@@ -19,6 +19,8 @@ from datetime import datetime
 from tqdm import tqdm
 from tensorflow.keras.callbacks import TensorBoard
 import collections
+import csv
+import time
 
 # Import necessary components from the pipeline
 # Use absolute imports assuming the package structure is correct
@@ -208,6 +210,8 @@ class SACTrainer:
         self.control_cfg = config.get('control', {})
         self.data_cfg = config['data']
 
+        # --- Cache useful config values ---
+        self.use_ternary = self.config.get('gru', {}).get('use_ternary', False)
         # --- Store PER config params --- #
         self.use_per = self.sac_cfg.get('use_per', False)
         self.per_alpha = self.sac_cfg.get('per_alpha', 0.6)
@@ -301,23 +305,85 @@ class SACTrainer:
         # 3. Load GRU Model
         model_path = os.path.join(gru_run_models_dir, f"gru_model_{gru_run_id}.keras")
         # Need a temporary GRU handler instance to use its load method
-        temp_gru_handler = GRUModelHandler(run_id="temp_load", models_dir="temp_load")
+        # Pass self.config (the trainer's config) to the GRUModelHandler
-        dependencies['gru_model'] = temp_gru_handler.load(model_path)
+        temp_gru_handler = GRUModelHandler(run_id="temp_load", models_dir="temp_load", config=self.config)
+        # Attempt to load the model using the handler
+        # The load method likely needs the full path, not just the directory and ID
+        loaded_model_info = temp_gru_handler.load(model_path=model_path) # Pass full path
+        
+        # Adjust based on what gru_handler.load returns
+        # Assuming it returns (model, info_dict) or None
+        if loaded_model_info and isinstance(loaded_model_info, tuple) and len(loaded_model_info) > 0:
+            dependencies['gru_model'] = loaded_model_info[0] # Get the actual model
             if dependencies['gru_model'] is None:
-             logger.error(f"Failed to load GRU model from {model_path}")
+                 logger.error(f"GRU handler load method returned None for model from {model_path}")
                  return None
+        elif loaded_model_info and not isinstance(loaded_model_info, tuple): # If it just returns the model
+             dependencies['gru_model'] = loaded_model_info
+        else: # If it returned None or unexpected tuple
+             logger.error(f"Failed to load GRU model using handler from {model_path}")
+             return None
+             
         logger.info(f"Loaded GRU model from {model_path}")
 
-        # 4. Load Optimal Temperature
+        # 4. Load Calibration Info and Parameters
-        temp_path = os.path.join(gru_run_models_dir, f"calibration_temp_{gru_run_id}.npy")
+        calib_info_path = os.path.join(gru_run_models_dir, f"calibration_info_{gru_run_id}.json")
+        calib_params = None
+        calib_method = None
         try:
-            dependencies['optimal_T'] = float(np.load(temp_path))
+            with open(calib_info_path, 'r') as f:
-            logger.info(f"Loaded optimal temperature T={dependencies['optimal_T']:.4f} from {temp_path}")
+                calib_info = json.load(f)
+            logger.info(f"Loaded calibration info: {calib_info}")
+            calib_method = calib_info.get("method")
+            params_filename = calib_info.get("params_filename")
+
+            if calib_method and params_filename:
+                params_path = os.path.join(gru_run_models_dir, params_filename)
+                if os.path.exists(params_path):
+                    if calib_method == "temperature":
+                        calib_params = float(np.load(params_path))
+                        logger.info(f"Loaded optimal temperature T={calib_params:.4f} from {params_path}")
+                    elif calib_method == "vector":
+                        # Load vector params (assuming saved as .npy)
+                        calib_params = np.load(params_path, allow_pickle=True)
+                        logger.info(f"Loaded vector calibration params from {params_path} (type: {type(calib_params)}, shape: {getattr(calib_params, 'shape', 'N/A')})")
+                    else:
+                        logger.warning(f"Unknown calibration method '{calib_method}' in info file. Cannot load params.")
+                else:
+                    logger.error(f"Calibration parameter file specified in info ({params_filename}) not found at {params_path}")
+                    return None # Fail if params file missing
+            elif calib_method in ["temperature_failed", "vector_failed", "skipped_mismatch", "temperature_error", "vector_error", "vector_unavailable"]:
+                logger.warning(f"Calibration was not successful or was skipped during GRU run (method: {calib_method}). SAC training may be suboptimal.")
+                # Decide if this is fatal. For now, let's allow it but store None.
+                calib_params = None
+            else:
+                logger.error(f"Invalid calibration info content: {calib_info}")
+                return None # Fail if info is invalid
+
+        except FileNotFoundError:
+            logger.error(f"Calibration info file not found at {calib_info_path}. Cannot determine calibration parameters.")
+            # Check for legacy temperature file for backward compatibility?
+            legacy_temp_path = os.path.join(gru_run_models_dir, f"calibration_temp_{gru_run_id}.npy")
+            if os.path.exists(legacy_temp_path):
+                 logger.warning("Calibration info file missing, but found legacy temperature file. Attempting to load it.")
+                 try:
+                     calib_params = float(np.load(legacy_temp_path))
+                     calib_method = "temperature" # Assume temperature
+                     logger.info(f"Loaded legacy optimal temperature T={calib_params:.4f} from {legacy_temp_path}")
+                 except Exception as legacy_e:
+                     logger.error(f"Failed to load legacy temperature file {legacy_temp_path}: {legacy_e}")
+                     return None # Fail if legacy load fails
+            else:
+                 logger.error("Neither calibration info file nor legacy temperature file found.")
+                 return None # Fail if no calibration info found
         except Exception as e:
-            logger.error(f"Failed to load optimal temperature from {temp_path}: {e}", exc_info=True)
+            logger.error(f"Failed to load calibration info/parameters: {e}", exc_info=True)
-            # Allow continuation without T? Or require it? Let's require it for now.
             return None
 
+        # Store method and params in dependencies
+        dependencies['calibration_method'] = calib_method
+        dependencies['calibration_params'] = calib_params
+            
         logger.info("--- Successfully loaded all GRU dependencies ---")
         return dependencies
 
@@ -349,107 +415,284 @@ class SACTrainer:
             df_raw.dropna(subset=['open', 'high', 'low', 'close', 'volume'], inplace=True)
             logger.info("Loaded raw data.")
 
-            # 2. Engineer Base Features (using a temporary FeatureEngineer)
+            # 2. Engineer Base Features (using config)
-            # Pass the *minimal* whitelist as a fallback if the loaded one causes issues
+            # FIX: Instantiate FeatureEngineer correctly using self.config
-            temp_feature_engineer = FeatureEngineer(minimal_whitelist=minimal_whitelist)
+            temp_feature_engineer = FeatureEngineer(config=self.config)
             df_engineered = temp_feature_engineer.add_base_features(df_raw)
             df_engineered.dropna(inplace=True) # Drop NaNs after feature eng
             if df_engineered.empty: raise ValueError("Dataframe empty after feature engineering.")
             logger.info("Engineered base features.")
 
-            # 3. Prune Features using *loaded* whitelist
+            # 3. Define Labels (on the full engineered data)
-            loaded_whitelist = gru_dependencies['whitelist']
-            missing_in_eng = [f for f in loaded_whitelist if f not in df_engineered.columns]
-            if missing_in_eng:
-                 raise ValueError(f"Features from loaded whitelist missing in engineered data: {missing_in_eng}")
-            df_features = df_engineered[loaded_whitelist]
-            logger.info("Pruned features using loaded whitelist.")
-
-            # 4. Define Labels
             horizon = self.config['gru'].get('prediction_horizon', 5)
             target_ret_col = f'fwd_log_ret_{horizon}'
-            target_dir_col = f'direction_label_{horizon}'
+            target_dir_col = f'direction_label3_{horizon}' if self.use_ternary else f'direction_label_{horizon}'
-            df_engineered[target_ret_col] = np.log(df_engineered['close'].shift(-horizon) / df_engineered['close'])
+            _EPS = 1e-9
+            if 'future_close' not in df_engineered.columns:
+                 df_engineered['future_close'] = df_engineered['close'].shift(-horizon)
+            if 'future_close' in df_engineered.columns:
+                 df_engineered[target_ret_col] = np.log(df_engineered['future_close'] / (df_engineered['close'] + _EPS))
+            else:
+                 df_engineered[target_ret_col] = np.log(df_engineered['close'].shift(-horizon) / (df_engineered['close'] + _EPS))
+
+            if self.use_ternary:
+                flat_sigma = self.config.get('gru', {}).get('flat_sigma_multiplier', 0.3)
+                if 'ATR_14' in df_engineered.columns:
+                    atr_aligned = df_engineered['ATR_14'].reindex(df_engineered.index).bfill().ffill()
+                    threshold = flat_sigma * atr_aligned / (df_engineered['close'] + _EPS)
+                    df_engineered[target_dir_col] = np.select(
+                        [df_engineered[target_ret_col] > threshold, df_engineered[target_ret_col] < -threshold],
+                        [2, 0], default=1
+                    )
+                else:
+                    raise ValueError("ATR_14 needed for ternary labels not found.")
+            else: # Binary case
                 df_engineered[target_dir_col] = (df_engineered[target_ret_col] > 0).astype(int)
-            # Align by dropping NaNs in targets AND ensuring indices match features
+
+            # 4. Align Features and Targets
             df_engineered.dropna(subset=[target_ret_col, target_dir_col], inplace=True)
-            common_index = df_features.index.intersection(df_engineered.index)
+            # Identify *all* potential feature columns (excluding targets)
-            if common_index.empty:
+            potential_feature_cols = [col for col in df_engineered.columns if col not in [target_ret_col, target_dir_col, 'future_close']]
-                raise ValueError("No common index between features and targets after label definition.")
+            # Ensure whitelist features are present in potential features
-            df_features = df_features.loc[common_index]
+            loaded_whitelist = gru_dependencies['whitelist']
-            df_targets = df_engineered.loc[common_index, [target_ret_col, target_dir_col]]
+            missing_whitelist_check = set(loaded_whitelist) - set(potential_feature_cols)
+            if missing_whitelist_check:
+                 logger.warning(f"Whitelist features missing from engineered columns: {missing_whitelist_check}. Adding with 0 fill.")
+                 for col in missing_whitelist_check:
+                      df_engineered[col] = 0.0
+                      potential_feature_cols.append(col)
+            # Now df_engineered contains all potential features and targets, aligned
             logger.info("Defined labels and aligned features/targets.")
 
-            # 5. Split Data (to get validation set indices)
+            # 5. Split Data (using aligned engineered data)
-            split_cfg = self.config['split_ratios']
+            wf_enabled = self.config.get('walk_forward', {}).get('enabled', False)
-            train_ratio, val_ratio = split_cfg['train'], split_cfg['validation']
+            if wf_enabled: logger.warning("Walk-forward enabled, but SAC uses split_ratios.")
-            total_len = len(df_features)
+            split_ratios = self.config.get('walk_forward', {}).get('split_ratios', {})
+            train_ratio = split_ratios.get('train', 0.6); val_ratio = split_ratios.get('validation', 0.2)
+            if not (0 < train_ratio < 1 and 0 < val_ratio < 1 and (train_ratio + val_ratio) <= 1.0):
+                 raise ValueError(f"Invalid split ratios: train={train_ratio}, validation={val_ratio}")
+            total_len = len(df_engineered)
             train_end_idx = int(total_len * train_ratio)
             val_end_idx = int(total_len * (train_ratio + val_ratio))
-            val_indices = df_features.index[train_end_idx:val_end_idx]
+            val_indices = df_engineered.index[train_end_idx:val_end_idx]
             if val_indices.empty: raise ValueError("Validation split resulted in empty indices.")
-            X_val_pruned = df_features.loc[val_indices]
+            df_val_aligned = df_engineered.loc[val_indices]
-            y_val = df_targets.loc[val_indices]
-            logger.info(f"Isolated validation set data (Features: {X_val_pruned.shape}, Targets: {y_val.shape}).")
 
-            # 6. Scale Validation Features using *loaded* scaler
+            # -- Determine columns expected by scaler --
             scaler = gru_dependencies['scaler']
-            numeric_cols = X_val_pruned.select_dtypes(include=np.number).columns
+            expected_scaler_features = []
-            X_val_scaled = X_val_pruned.copy()
+            if hasattr(scaler, 'feature_names_in_'):
-            if not numeric_cols.empty:
+                expected_scaler_features = scaler.feature_names_in_.tolist()
-                 X_val_scaled[numeric_cols] = scaler.transform(X_val_pruned[numeric_cols])
+                logger.debug(f"Scaler expects features: {expected_scaler_features}")
-            logger.info("Scaled validation features using loaded scaler.")
+            else:
+                # Fallback: Use all numeric columns if scaler has no names
+                logger.warning("Scaler has no feature names saved. Assuming it was fit on all numeric columns present at that time.")
+                # This is risky; need to ensure columns match implicitly
+                expected_scaler_features = df_val_aligned.select_dtypes(include=np.number).columns.tolist()
+                # Manually exclude known target columns if needed as a safeguard
+                expected_scaler_features = [f for f in expected_scaler_features if f not in [target_ret_col, target_dir_col]]
+                logger.debug(f"Falling back to using numeric columns for scaler: {expected_scaler_features}")
 
-            # 7. Create Validation Sequences
+            # Ensure all expected scaler features exist in the validation slice
+            missing_for_scaler = set(expected_scaler_features) - set(df_val_aligned.columns)
+            if missing_for_scaler:
+                # Attempt to add missing with 0 fill (e.g., if a feature wasn't calculable for val split)
+                logger.warning(f"Columns expected by scaler missing from validation data: {missing_for_scaler}. Adding with 0 fill.")
+                for col in missing_for_scaler:
+                    df_val_aligned[col] = 0.0
+                # Re-verify
+                missing_for_scaler = set(expected_scaler_features) - set(df_val_aligned.columns)
+                if missing_for_scaler:
+                    raise ValueError(f"Could not prepare all features expected by scaler, still missing: {missing_for_scaler}")
+
+            # Isolate the exact features needed for the scaler transform
+            X_val_engineered_for_scaler = df_val_aligned[expected_scaler_features]
+            y_val = df_val_aligned[[target_ret_col, target_dir_col]]
+            logger.info(f"Isolated validation set data for scaler (Features: {X_val_engineered_for_scaler.shape}, Targets: {y_val.shape}).")
+
+            # 6. Scale the features expected by the scaler
+            X_val_scaled_full = X_val_engineered_for_scaler.copy()
+            try:
+                X_val_scaled_full[expected_scaler_features] = scaler.transform(X_val_engineered_for_scaler)
+                logger.info("Scaled validation features using loaded scaler.")
+            except ValueError as e:
+                 logger.error(f"Error applying scaler transform even after aligning columns: {e}")
+                 raise
+
+            # 7. WORKAROUND: Drop non-feature columns (like future_close) *after* scaling
+            # Identify actual features intended for the model (whitelist)
+            loaded_whitelist = gru_dependencies['whitelist']
+            cols_to_keep_after_scale = [f for f in loaded_whitelist if f in X_val_scaled_full.columns]
+            missing_whitelist_post_scale = set(loaded_whitelist) - set(cols_to_keep_after_scale)
+            if missing_whitelist_post_scale:
+                 # This shouldn't happen if whitelist features were in expected_scaler_features
+                 logger.warning(f"Whitelist features missing AFTER scaling: {missing_whitelist_post_scale}. This might indicate issues.")
+
+            if not cols_to_keep_after_scale:
+                 raise ValueError("Whitelist resulted in no features remaining after scaling step.")
+
+            X_val_scaled_eng = X_val_scaled_full[cols_to_keep_after_scale].copy()
+            logger.info(f"Selected whitelisted features after scaling. Shape: {X_val_scaled_eng.shape}")
+
+            # 8. Pruning step is now effectively done by selecting whitelist columns above.
+            X_val_pruned_scaled = X_val_scaled_eng # Rename for clarity in subsequent steps
+            logger.info(f"Using whitelisted scaled features for sequencing. Final shape: {X_val_pruned_scaled.shape}")
+
+            # 9. Create Validation Sequences (using the pruned & scaled data)
             lookback = self.config['gru']['lookback']
-            X_val_seq = []
+            X_val_seq_list, y_val_seq_targets_list, val_seq_indices = [], [], []
-            y_val_seq_targets = [] # Store corresponding targets
+            features_np_arr = X_val_pruned_scaled.values
-            val_seq_indices = []   # Store corresponding indices
+            targets_np_arr = y_val.values
-            features_np = X_val_scaled.values
+            for i in range(lookback, len(features_np_arr)):
-            targets_np = y_val.values # Contains both ret and dir
+                X_val_seq_list.append(features_np_arr[i-lookback : i])
-            for i in range(lookback, len(features_np)):
+                y_val_seq_targets_list.append(targets_np_arr[i])
-                X_val_seq.append(features_np[i-lookback : i])
-                y_val_seq_targets.append(targets_np[i]) # Target corresponds to end of sequence
                 val_seq_indices.append(y_val.index[i])
 
-            if not X_val_seq:
+            if not X_val_seq_list: raise ValueError("Validation sequence creation resulted in empty list.")
-                 raise ValueError("Validation sequence creation resulted in empty list.")
+            X_val_seq = np.array(X_val_seq_list)
-                 
+            y_val_seq_targets = np.array(y_val_seq_targets_list)
-            X_val_seq = np.array(X_val_seq)
+            actual_ret_val_seq = y_val_seq_targets[:, 0]
-            y_val_seq_targets = np.array(y_val_seq_targets)
+            y_dir_val_seq = y_val_seq_targets[:, 1]
-            actual_ret_val_seq = y_val_seq_targets[:, 0] # First column is return
-            y_dir_val_seq = y_val_seq_targets[:, 1] # Second column is direction
             logger.info(f"Created validation sequences (X shape: {X_val_seq.shape}).")
 
-            # 8. Get GRU Predictions on Validation Sequences using *loaded* GRU model
+            # 10. Get GRU Predictions using loaded GRU model directly
             gru_model = gru_dependencies['gru_model']
-            # Use a temporary handler instance with the loaded model
+            logger.info(f"Generating GRU predictions using loaded model (type: {type(gru_model)}).")
-            temp_gru_handler = GRUModelHandler(run_id="temp_predict", models_dir="temp")
+
-            temp_gru_handler.model = gru_model # Assign the loaded model
+            # Check model type and predict accordingly
-            predictions_val = temp_gru_handler.predict(X_val_seq)
+            if not hasattr(gru_model, 'predict'):
-            if predictions_val is None or len(predictions_val) < 3:
+                 raise TypeError("Loaded GRU model object does not have a 'predict' method.")
-                 raise ValueError("GRU prediction on validation sequences failed.")
+
-            mu_val_pred = predictions_val[0].flatten()
+            predictions_val = gru_model.predict(X_val_seq)
-            log_sigma_val_pred = predictions_val[1][:, 1].flatten()
+
-            p_raw_val_pred = predictions_val[2].flatten()
+            # --- Explicit Shape Logging --- #
-            sigma_val_pred = np.exp(log_sigma_val_pred)
+            if isinstance(predictions_val, list):
-            logger.info("Generated GRU predictions on validation sequences.")
+                logger.info(f"DEBUG: GRU prediction output list length: {len(predictions_val)}")
+                if len(predictions_val) > 0: logger.info(f"DEBUG: Shape of predictions_val[0]: {predictions_val[0].shape}")
+                if len(predictions_val) > 1: logger.info(f"DEBUG: Shape of predictions_val[1]: {predictions_val[1].shape}")
+                if len(predictions_val) > 2: logger.info(f"DEBUG: Shape of predictions_val[2]: {predictions_val[2].shape}")
+            else:
+                logger.info(f"DEBUG: GRU prediction output type: {type(predictions_val)}")
+            # --- End Explicit Shape Logging --- #
+
+            # --- Infer output structure (Corrected Indexing) ---
+            mu_val_pred, sigma_val_pred, p_raw_val_pred, logits_val_pred = None, None, None, None
+
+            if isinstance(predictions_val, list) and len(predictions_val) == 2:
+                logger.debug(f"GRU model returned list of {len(predictions_val)} outputs. Correcting index assumption: [mu, dir].")
+                # Corrected Indexing based on logs:
+                mu_output = predictions_val[0] # Index 0 has shape (N, 1)
+                dir_output = predictions_val[1] # Index 1 has shape (N, 3)
+
+                # Extract Mu
+                if mu_output.ndim == 2 and mu_output.shape[-1] == 1:
+                     mu_val_pred = mu_output.flatten()
+                else:
+                     # Corrected error message to reflect index 0 check
+                     raise ValueError(f"Unexpected shape for mu output (index 0): {mu_output.shape}. Expected (N, 1).")
+
+                # Handle missing Sigma - Use a default/fallback
+                logger.warning("Log_sigma_sq output not found in model prediction. Using default sigma=0.1.")
+                sigma_val_pred = np.ones_like(mu_val_pred) * 0.1
+
+                # Determine direction output type from dir_output (index 1)
+                if self.use_ternary:
+                    if dir_output.ndim == 2 and dir_output.shape[-1] == 3:
+                        # Assume dir_output (index 1) is logits if ternary
+                        logits_val_pred = dir_output
+                        from scipy.special import softmax # Local import ok here
+                        p_raw_val_pred = softmax(logits_val_pred, axis=-1)
+                        logger.info("Inferred ternary output (logits assumed at index 1). Calculated raw probabilities.")
+                    else:
+                         raise ValueError(f"Expected ternary output shape (N, 3) at index 1, got {dir_output.shape}")
+                else: # Binary
+                    # Binary case is tricky if dir_output (index 1) is still (N, 3)
+                    # This shouldn't happen if the model structure changes for binary
+                    # Let's assume the mu_output (index 0) might be P(up) in binary?
+                    logger.warning(f"Binary mode, dir_output shape is {dir_output.shape}. Trying to use mu output (index 0) as P(up) - THIS IS RISKY.")
+                    if mu_output.ndim == 2 and mu_output.shape[-1] == 1:
+                         p_raw_val_pred = mu_output.flatten()
+                         epsilon = 1e-7
+                         p_clipped = np.clip(p_raw_val_pred, epsilon, 1 - epsilon)
+                         logits_val_pred = np.log(p_clipped / (1 - p_clipped))
+                         logger.info("Using mu output (index 0) as P(up) for binary mode. Inferred raw logits.")
+                    else:
+                         raise ValueError(f"Cannot determine binary P(up) output. mu_output shape: {mu_output.shape}")
+            else:
+                 raise ValueError(f"GRU prediction failed or returned unexpected type/structure: {type(predictions_val)}, length: {len(predictions_val) if isinstance(predictions_val, list) else 'N/A'}. Expected list of 2 arrays.")
+
+            # Final check for necessary predictions
+            if mu_val_pred is None or sigma_val_pred is None or p_raw_val_pred is None:
+                 raise ValueError("Failed to extract mu, sigma, or raw probabilities from GRU model output.")
 
             # Verify lengths
             n_seq = len(X_val_seq)
             if not (len(mu_val_pred) == n_seq and len(sigma_val_pred) == n_seq and \
-                      len(p_raw_val_pred) == n_seq and len(actual_ret_val_seq) == n_seq):
+                    p_raw_val_pred.shape[0] == n_seq and len(actual_ret_val_seq) == n_seq):
-                 raise ValueError(f"Length mismatch after validation predictions: Expected {n_seq}, got mu={len(mu_val_pred)}, sigma={len(sigma_val_pred)}, p_raw={len(p_raw_val_pred)}, ret={len(actual_ret_val_seq)}")
+                 raise ValueError(f"Length mismatch after validation predictions: Expected {n_seq}, "
+                                  f"got mu={len(mu_val_pred)}, sigma={len(sigma_val_pred)}, "
+                                  f"p_raw={p_raw_val_pred.shape[0]}, ret={len(actual_ret_val_seq)}")
 
-            # 9. Calibrate Predictions using *loaded* optimal_T
+            # 11. Calibrate Predictions using loaded parameters
-            optimal_T = gru_dependencies['optimal_T']
+            calib_method = gru_dependencies.get('calibration_method')
-            # Use a temporary calibrator instance
+            calib_params = gru_dependencies.get('calibration_params')
-            temp_calibrator = Calibrator(edge_threshold=0.5) # Edge threshold doesn't matter here
+            p_cal_val_pred = None
-            temp_calibrator.optimal_T = optimal_T
-            p_cal_val_pred = temp_calibrator.calibrate(p_raw_val_pred)
-            logger.info(f"Calibrated validation predictions using loaded T={optimal_T:.4f}.")
 
-            # 10. Return the necessary components for the TradingEnv
+            if calib_method == "temperature" and calib_params is not None:
+                optimal_T = calib_params
+                # Temperature scaling needs logits
+                if logits_val_pred is not None:
+                     scaled_logits = logits_val_pred / optimal_T
+                     if self.use_ternary:
+                         p_cal_val_pred = softmax(scaled_logits, axis=-1)
+                         logger.info(f"Applied temperature scaling (T={optimal_T:.4f}) to ternary logits.")
+                     else: # Binary - apply sigmoid
+                         p_cal_val_pred = 1 / (1 + np.exp(-scaled_logits)) # Sigmoid
+                         logger.info(f"Applied temperature scaling (T={optimal_T:.4f}) to binary logits.")
+                else:
+                     logger.error(f"Cannot apply temperature scaling (method={calib_method}): Raw logits not available.")
+                     p_cal_val_pred = p_raw_val_pred # Fallback to raw
+
+            elif calib_method == "vector" and calib_params is not None:
+                # Vector calibration needs logits
+                if logits_val_pred is not None:
+                    try:
+                        from gru_sac_predictor.src.calibrator_vector import VectorCalibrator # Ensure import
+                        temp_vector_calibrator = VectorCalibrator()
+                        temp_vector_calibrator.optimal_params = calib_params # Set loaded params
+                        # Calibrate expects logits, returns probabilities
+                        p_cal_val_pred = temp_vector_calibrator.calibrate(logits_val_pred)
+                        logger.info(f"Calibrated validation predictions using loaded vector parameters.")
+                    except ImportError:
+                        logger.error("Cannot import VectorCalibrator. Vector calibration step skipped.")
+                        p_cal_val_pred = p_raw_val_pred # Fallback to raw
+                    except Exception as vec_e:
+                        logger.error(f"Error during vector calibration application: {vec_e}", exc_info=True)
+                        p_cal_val_pred = p_raw_val_pred # Fallback to raw
+                else:
+                    logger.error(f"Cannot apply vector calibration (method={calib_method}): Raw logits not available.")
+                    p_cal_val_pred = p_raw_val_pred # Fallback to raw
+
+            else: # No calibration or failed
+                logger.warning(f"Calibration method was '{calib_method}' or params were None or previous step failed. Using RAW predictions for SAC environment.")
+                p_cal_val_pred = p_raw_val_pred # Use the raw probs (softmaxed for ternary, P(up) for binary)
+
+            # Final check for calibrated probabilities
+            if p_cal_val_pred is None:
+                 logger.error("Failed to obtain final (calibrated or raw) predictions. Cannot proceed.")
+                 return None
+            # Verify shape of final probabilities
+            expected_final_dim = 3 if self.use_ternary else 1
+            if p_cal_val_pred.ndim == 2 and p_cal_val_pred.shape[-1] == expected_final_dim:
+                 if not self.use_ternary: p_cal_val_pred = p_cal_val_pred.flatten() # Flatten binary case
+            elif p_cal_val_pred.ndim == 1 and not self.use_ternary and expected_final_dim == 1:
+                 pass # Already flat for binary
+            else:
+                 logger.error(f"Final probability array has unexpected shape: {p_cal_val_pred.shape}. Expected dim {expected_final_dim}.")
+                 return None
+
+
+            # 12. Return the necessary components for the TradingEnv
             logger.info("--- Successfully prepared validation data for SAC Environment ---")
             return mu_val_pred, sigma_val_pred, p_cal_val_pred, actual_ret_val_seq
 
@@ -512,358 +755,206 @@ class SACTrainer:
             logger.warning(f"SAC agent path not found for resume: {load_path}. Starting fresh.")
             
     def _training_loop(self, agent: SACTradingAgent, env: TradingEnv) -> str | None:
-        """Runs the main SAC training loop."""
+        """The main SAC training loop."""
-        buffer_max_size = self.sac_cfg.get('buffer_max_size', 100000)
+        total_steps = self.sac_cfg.get('total_training_steps', 100000)
-        min_buffer_size = self.sac_cfg.get('min_buffer_size', 10000)
+        start_steps = self.sac_cfg.get('start_steps', 10000)
+        update_after = self.sac_cfg.get('update_after', 1000)
+        update_every = self.sac_cfg.get('update_every', 50)
+        save_freq = self.sac_cfg.get('save_freq', 5000)
+        log_freq = self.sac_cfg.get('log_freq', 100)
+        buffer_capacity = self.sac_cfg.get('buffer_capacity', 1000000)
         batch_size = self.sac_cfg.get('batch_size', 256)
-        total_training_steps = self.sac_cfg.get('total_training_steps', 100000)
 
-        # --- Initialize Replay Buffer (Potentially PER) --- #
+        # Initialize Replay Buffer (Standard or Prioritized)
         if self.use_per:
-            logger.info(f"Initializing Prioritized Replay Buffer (Capacity={buffer_max_size}, alpha={self.per_alpha}, beta_start={self.per_beta_start}, beta_frames={self.per_beta_frames})")
+            logger.info(f"Using Prioritized Replay Buffer (Capacity: {buffer_capacity})")
             replay_buffer = PrioritizedReplayBuffer(
-                buffer_max_size, 
+                capacity=buffer_capacity, 
-                alpha=self.per_alpha, 
+                alpha=self.per_alpha, # Initial alpha 
                 beta_start=self.per_beta_start,
                 beta_frames=self.per_beta_frames
             )
         else:
-            logger.info(f"Initializing Standard Replay Buffer (Deque, Capacity={buffer_max_size})")
+            logger.info(f"Using Standard Replay Buffer (Capacity: {buffer_capacity})")
-            replay_buffer = collections.deque(maxlen=buffer_max_size)
+            replay_buffer = collections.deque(maxlen=buffer_capacity)
-            replay_buffer.counter = 0 # Add counter for uniform sampling logic
-        # --- End Buffer Init --- #
 
-        # --- Oracle Seeding (Revision 4-B) --- #
+        # TensorBoard setup
-        oracle_seeding_pct = self.sac_cfg.get('oracle_seeding_pct', 0.0)
+        tb_callback = TensorBoard(log_dir=self.sac_tb_log_dir)
-        num_existing_samples = len(replay_buffer) # Count samples potentially loaded during resume
+        # --- Revision 4: Set model for TensorBoard --- #
-        target_seed_steps = int(buffer_max_size * oracle_seeding_pct)
+        # Check if agent has actor/critic models accessible
-        actual_seed_steps = max(0, target_seed_steps - num_existing_samples)
+        # This depends heavily on SACTradingAgent implementation
-        
+        # Assuming agent.actor and agent.critic1/2 are the models
-        if actual_seed_steps > 0:
+        if hasattr(agent, 'actor') and hasattr(agent, 'critic1') and hasattr(agent, 'critic2'):
-            logger.info(f"Performing Oracle Seeding: Adding ~{actual_seed_steps} steps ({oracle_seeding_pct * 100:.1f}% target) with heuristic policy...")
+             tb_callback.set_model(agent.actor) # Link to one model is often sufficient 
-            
+             logger.info("TensorBoard callback linked to SAC agent model (actor).")
-            # Need edge threshold from config (used in the heuristic action)
-            edge_threshold_heuristic = self.config.get('calibration', {}).get('edge_threshold')
-            if edge_threshold_heuristic is None:
-                logger.error("Cannot perform oracle seeding: 'calibration.edge_threshold' not found in config.")
-            elif edge_threshold_heuristic <= 0:
-                 logger.warning(f"Edge threshold for heuristic is {edge_threshold_heuristic:.3f}. Oracle seeding action may be ill-defined. Using 0.01 instead.")
-                 edge_threshold_heuristic = 0.01
         else:
-                state = env.reset() # Start seeding from the beginning of the env data
+             logger.warning("Could not link TensorBoard callback to agent models (actor/critic not found).")
-                n_seeded = 0
+        # --- End Revision 4 ---
+
+        # --- Initialize optional imputed transition logger --- #
+        imputed_log_path = os.path.join(self.sac_run_results_dir, 'sac_imputed_transitions.csv')
+        imputed_log_file = None
+        imputed_csv_writer = None
         try:
-                    for _ in tqdm(range(actual_seed_steps), desc="Oracle Seeding", file=sys.stdout, leave=False):
+            imputed_log_file = open(imputed_log_path, 'w', newline='')
-                        if len(replay_buffer) >= buffer_max_size: 
+            imputed_csv_writer = csv.writer(imputed_log_file)
-                             logger.warning("Buffer full during oracle seeding.")
+            imputed_csv_writer.writerow(['step', 'imputed_handling_mode', 'action', 'reward', 'position_before', 'position_after'])
-                             break 
+            logger.info(f"Logging imputed transitions details to: {imputed_log_path}")
+        except IOError as e:
+            logger.error(f"Failed to open imputed transition log file {imputed_log_path}: {e}. Logging disabled.")
+            if imputed_log_file: imputed_log_file.close()
+            imputed_log_file = None
+            imputed_csv_writer = None
+        # --- End init imputed logger --- #
 
-                        current_state_for_seeding = state # Keep original state before potential filtering
-                        if self.state_filter:
-                             state = self.state_filter(state, update=False) # Apply filter, don't update filter stats during seeding
-                         
-                        # Heuristic action based on edge
-                        # State: [mu, sigma, edge, |mu|/sigma, position] (assuming this order)
-                        try:
-                             # Ensure state has enough elements before indexing
-                             if len(state) < 3:
-                                  logger.error(f"State vector too short ({len(state)} elements) for oracle seeding. Stopping seeding.")
-                                  break
-                             edge = state[2] # Assuming edge is the 3rd element (index 2)
-                        except IndexError:
-                             logger.error(f"IndexError accessing edge (state[2]) during oracle seeding. State shape: {state.shape if hasattr(state, 'shape') else type(state)}. Stopping seeding.")
-                             break
-                             
-                        oracle_action = np.clip(edge / edge_threshold_heuristic, -1.0, 1.0)
-                        
-                        next_state, reward, done, _ = env.step(oracle_action)
-                        
-                        # Store experience with the state *before* filtering (as agent expects raw state)
-                        experience = (current_state_for_seeding, oracle_action, reward, next_state, done)
-                        if self.use_per:
-                             # Mark as seeded
-                             replay_buffer.add(replay_buffer.max_priority, experience, is_seeded=True)
-                        elif hasattr(replay_buffer, 'counter'):
-                             replay_buffer.append(experience)
-                             replay_buffer.counter += 1
-                        n_seeded += 1
-                        
-                        state = next_state
-                        if done:
-                            # Reset env if it finishes during seeding, continue seeding if needed
-                            logger.info("Environment finished during oracle seeding. Resetting.")
         state = env.reset()
-                    logger.info(f"Oracle seeding completed. Added {n_seeded} experiences.")
+        # Normalize initial state if filter is active
-                except Exception as seed_err:
-                     logger.error(f"Error during oracle seeding loop: {seed_err}. Proceeding with {n_seeded} seeded samples.", exc_info=True)
-                     
-        else:
-            logger.info("Oracle seeding skipped (percentage=0 or buffer already seeded enough). Existing samples: {num_existing_samples}")
-        # --- End Oracle Seeding --- #
-        
-        # --- Training Loop Setup --- #
-        start_learning_after_steps = self.sac_cfg.get('start_learning_after_steps', 1000)
-        save_checkpoint_freq = self.sac_cfg.get('save_checkpoint_freq_steps', 50000)
-        total_steps = self.sac_cfg.get('total_training_steps', 100000)
-        
-        # Revision 5: Alpha annealing params
-        alpha_start = self.sac_cfg.get('per_alpha_start', 0.6)
-        alpha_end = self.sac_cfg.get('per_alpha_end', 0.4)
-        current_alpha = alpha_start # Initialize alpha
-
-        summary_writer = tf.summary.create_file_writer(self.sac_tb_log_dir)
-        episode_reward = 0
-        episode_steps = 0
-        episode_rewards_log = [] # For saving reward history
-        best_eval_score = -np.inf # Placeholder for potential periodic evaluation
-        final_agent_path = None
-        
-        logger.info(f"Starting SAC training loop for {total_training_steps} steps...")
-        logger.info(f" Min buffer size: {min_buffer_size}, Batch size: {batch_size}, Updates/step: {updates_per_step}")
-        logger.info(f" Log interval: {log_interval}, Checkpoint interval: {checkpoint_interval}")
-        
-        summary_writer = tf.summary.create_file_writer(self.sac_tb_log_dir)
-        pbar = tqdm(range(1, total_training_steps + 1), desc="SAC Training", file=sys.stdout)
-        last_saved_path = None # Track last successful save
-
-        # --- Main Training Loop --- #
-        for step in pbar:
-            # --- Revision 5: Update Annealed Alpha --- #
-            if self.use_per:
-                alpha_fraction = min(1.0, step / total_steps)
-                current_alpha = alpha_start + alpha_fraction * (alpha_end - alpha_start)
-                # Calculate Seed Decay factor for IS weights
-                seed_decay_factor = max(0.0, 1.0 - step / self.per_seed_decay_steps)
-            # --- End Revision 5 --- #
-
-            original_state = state # Keep original state for buffer
         if self.state_filter:
-                 state = self.state_filter(state, update=True) # Apply filter and update running stats
+            state = self.state_filter(state, update=True) # Update filter with initial state
 
-            if step < start_learning_after_steps:
+        total_reward = 0.0
-                # Use random actions during warmup to explore
+        start_time = time.time()
-                action = np.random.uniform(env.action_space.low, env.action_space.high, size=env.action_space.shape)
+
+        logger.info(f"Starting SAC training loop for {total_steps} steps...")
+        for step in tqdm(range(total_steps), desc="SAC Training Steps"):
+            # Store state before taking action (for replay buffer)
+            state_before_action = state 
+            position_before_action = env.current_position # Store for imputed log
+
+            # Select action
+            if step < start_steps:
+                action = env.action_space.sample()[0] # Sample random action
             else:
-                # Get action from agent
+                action = agent.select_action(state)
-                action, _ = agent.select_action(state)
 
-            # Environment step
+            # Step the environment
-            next_state, reward, done, info = env.step(action[0]) # Env expects single float action
+            next_state_raw, reward, done, info = env.step(action)
 
-            # Store experience in buffer (use original state)
+            # Check if the step was skipped due to imputed bar handling
-            experience = (original_state, action, reward, next_state, done)
+            is_skipped = info.get('is_imputed_step_skipped', False)
+
+            # Normalize next state if filter is active
+            if self.state_filter:
+                 next_state = self.state_filter(next_state_raw, update=True) # Update filter
+            else:
+                 next_state = next_state_raw
+
+            # Store transition ONLY if the step was NOT skipped
+            if not is_skipped:
                 if self.use_per:
-                 # Mark as seeded
+                    # Add with initial error=1 (or max priority) and is_seeded=False
-                 replay_buffer.add(replay_buffer.max_priority, experience, is_seeded=True)
+                    # The error will be updated after the first training step on this sample.
-            elif hasattr(replay_buffer, 'counter'): 
+                    replay_buffer.add(error=1.0, sample=(state_before_action, action, reward, next_state, done))
-                 replay_buffer.append(experience)
+                else:
-                 replay_buffer.counter += 1 # Increment counter for deque
+                    replay_buffer.append((state_before_action, action, reward, next_state, done))
+            else:
+                 logger.debug(f"Step {step}: Skipped adding transition to buffer due to imputed bar handling (mode=skip).")
+
+            # --- Log imputed transition details to CSV --- #
+            # Check if the *previous* step was imputed and handled by hold/penalty
+            # env.current_step was incremented *inside* env.step()
+            was_imputed_idx = env.current_step - 1
+            if 0 <= was_imputed_idx < env.n_steps and env.bar_imputed[was_imputed_idx] and not is_skipped and imputed_csv_writer:
+                 # Don't log if skipped, log if hold/penalty was applied
+                 imputed_handling_mode = self.config.sac.get('imputed_handling', 'unknown')
+                 action_taken = env.current_position if imputed_handling_mode == 'hold' else action # Action applied or agent's intended action
+                 position_after_action = env.current_position
+                 try:
+                     imputed_csv_writer.writerow([
+                         was_imputed_idx, # Log the actual step number where imputation occurred
+                         imputed_handling_mode,
+                         f"{action_taken:.4f}",
+                         f"{reward:.6f}", # Log the reward received for this step
+                         f"{position_before_action:.4f}",
+                         f"{position_after_action:.4f}"
+                     ])
+                 except Exception as log_e:
+                     logger.warning(f"Failed to write imputed transition to CSV: {log_e}")
+            # --- End Log imputed transition --- #
 
             state = next_state
-            current_episode_reward += reward
+            total_reward += reward
-            current_episode_steps += 1
 
-            # Perform SAC updates
+            # Perform SAC agent updates
-            if len(replay_buffer) >= min_buffer_size:
+            if step >= update_after and step % update_every == 0:
-                for _ in range(updates_per_step):
+                for update_i in range(update_every): # Perform multiple updates per interval
-                    sample_indices, batch_with_seed_flags, importance_weights = None, None, None # Initialize
-                    # --- Sample from Trainer's Buffer --- #
                     if self.use_per:
-                         sample_indices, batch_with_seed_flags, importance_weights = replay_buffer.sample(batch_size)
+                        if len(replay_buffer) > batch_size:
-                         # --- Revision 5: Apply Seed Decay to IS Weights --- #
+                            idxs, batch_data, is_weights = replay_buffer.sample(batch_size, beta=replay_buffer.beta) # Use annealed beta
-                         batch = [] # Store only the samples
+                            # Unpack batch_data which contains (sample, is_seeded) tuples
-                         seed_mask = np.zeros_like(importance_weights, dtype=bool)
+                            batch = [item[0] for item in batch_data]
-                         for i, (sample, is_seeded) in enumerate(batch_with_seed_flags):
+                            update_info = agent.update(batch, is_weights=is_weights, per_beta=replay_buffer.beta)
-                             batch.append(sample)
+                            if update_info and 'td_errors' in update_info:
-                             if is_seeded:
+                                 # Update priorities using current annealed alpha
-                                 seed_mask[i] = True
+                                 current_alpha = agent.get_current_per_alpha(step)
-                         # Apply decay factor to IS weights of seeded samples
+                                 replay_buffer.update_priorities(idxs, update_info['td_errors'], alpha=current_alpha)
-                         # --- Revision 6: Log IS weight decay --- #
-                         num_seeded_in_batch = np.sum(seed_mask)
-                         if num_seeded_in_batch > 0:
-                             original_weights_seeded = importance_weights[seed_mask].copy()
-                             importance_weights[seed_mask] *= seed_decay_factor
-                             weights_after_decay = importance_weights[seed_mask]
-                             # Log periodically or if decay factor changes significantly
-                             if step % self.log_interval == 0: # Align with other logging
-                                 logger.info(f"Step {step}: Applied IS decay factor {seed_decay_factor:.4f} to {num_seeded_in_batch} seeded samples in batch.")
-                                 # Optional: Log weight changes
-                                 # logger.debug(f"    Weights before: {np.round(original_weights_seeded, 3)}")
-                                 # logger.debug(f"    Weights after: {np.round(weights_after_decay, 3)}")
-                         # --- End Revision 6 --- #
-                         importance_weights_tensor = tf.convert_to_tensor(importance_weights, dtype=tf.float32)
-                    else: # Uniform sampling from deque
-                         current_buffer_size = min(replay_buffer.counter, replay_buffer.maxlen)
-                         if current_buffer_size < batch_size:
-                              sample_indices = np.random.choice(current_buffer_size, batch_size, replace=True)
                         else:
-                              sample_indices = np.random.choice(current_buffer_size, batch_size, replace=False)
+                            continue # Not enough samples yet for PER
-                         batch = [replay_buffer[i] for i in sample_indices]
+                    else: # Standard buffer
-                         importance_weights_tensor = None # No IS weights for uniform
+                        if len(replay_buffer) > batch_size:
-                    # --- End Sampling --- #
+                            indices = np.random.choice(len(replay_buffer), size=batch_size, replace=False)
-                         
+                            batch = [replay_buffer[i] for i in indices]
-                    state_batch, action_batch, reward_batch, next_state_batch, done_batch = map(np.stack, zip(*batch))
+                            update_info = agent.update(batch)
-                    
-                    # Apply state filter to sampled states if enabled
-                    if self.state_filter:
-                         state_batch_filtered = self.state_filter(state_batch, update=False) 
-                         next_state_batch_filtered = self.state_filter(next_state_batch, update=False)
                         else:
-                         state_batch_filtered = state_batch
+                             continue # Not enough samples yet
-                         next_state_batch_filtered = next_state_batch
 
-                    # Convert batch to tensors
+                    # Log training metrics (losses, Q-values, alpha) to TensorBoard
-                    state_tensor = tf.convert_to_tensor(state_batch_filtered, dtype=tf.float32)
+                    if update_info and step % log_freq == 0 and update_i == 0: # Log once per interval
-                    action_tensor = tf.convert_to_tensor(action_batch, dtype=tf.float32)
+                        with tb_callback.writer.as_default():
-                    reward_tensor = tf.convert_to_tensor(reward_batch, dtype=tf.float32)
+                            for key, value in update_info.items():
-                    next_state_tensor = tf.convert_to_tensor(next_state_batch_filtered, dtype=tf.float32)
+                                if key != 'td_errors': # Don't log TD errors directly
-                    done_tensor = tf.convert_to_tensor(done_batch, dtype=tf.float32)
+                                     tf.summary.scalar(f'sac/{key}', value, step=step)
+                        # logger.debug(f"Step {step}: Logged SAC metrics to TensorBoard.")
 
-                    # --- Call agent's train method with the batch --- #
+            # Check environment done state
-                    loss_info = agent.train(
-                        state_tensor, 
-                        action_tensor, 
-                        reward_tensor, 
-                        next_state_tensor, 
-                        done_tensor,
-                        importance_weights=importance_weights_tensor # Pass potentially decayed weights for PER
-                    )
-                    # --- End Agent Update Call --- #
-                    
-                    # --- Revision 5: Update PER priorities with current alpha --- #
-                    if self.use_per and loss_info is not None:
-                        td_errors = loss_info.get('td_errors') # Get TD errors from loss_info dict
-                        if td_errors is not None:
-                             replay_buffer.update_priorities(sample_indices, td_errors, current_alpha)
-                             # --- Revision 6: Log TD Error Distribution --- #
-                             # Convert to numpy if it's a tensor
-                             if tf.is_tensor(td_errors):
-                                 td_errors = td_errors.numpy()
-                             td_error_abs = np.abs(td_errors)
-                             log_hist_interval = self.sac_cfg.get('log_hist_interval', 5000)
-                             # Log percentiles every step where update happens
-                             step_info_for_logging = {
-                                 'td_error_abs_p25': np.percentile(td_error_abs, 25),
-                                 'td_error_abs_p50': np.percentile(td_error_abs, 50),
-                                 'td_error_abs_p75': np.percentile(td_error_abs, 75),
-                                 'td_error_abs_p95': np.percentile(td_error_abs, 95),
-                                 'td_error_abs_max': np.max(td_error_abs),
-                                 'td_error_abs_mean': np.mean(td_error_abs)
-                             }
-                             # Log histogram to TensorBoard periodically
-                             if log_hist_interval > 0 and step % log_hist_interval == 0:
-                                 try:
-                                     with summary_writer.as_default(step=step):
-                                         tf.summary.histogram('td_error_abs_distribution', td_error_abs, description='Absolute TD Errors')
-                                     summary_writer.flush()
-                                     logger.debug(f"Step {step}: Logged TD error histogram to TensorBoard.")
-                                 except Exception as hist_err:
-                                     logger.warning(f"Failed to log TD error histogram to TensorBoard at step {step}: {hist_err}")
-                             # --- End Revision 6 --- #
-                        else: 
-                             logger.warning(f"Step {step}: TD errors not found in loss_info from agent.train(). Cannot update PER priorities or log distribution.")
-                    # --- End Revision 5 --- #
-
-            # Logging and Checkpointing
-            if step % log_interval == 0 and len(replay_buffer) >= min_buffer_size:
-                with summary_writer.as_default(): # Use context manager
-                    # Log loss_info if available
-                    if 'loss_info' in locals() and loss_info:
-                         for k, v in loss_info.items():
-                              # Skip logging raw TD errors tensor here
-                              if k != 'td_errors':
-                                   tf.summary.scalar(f'loss/{k}', v, step=step)
-                         tf.summary.scalar('alpha', agent.log_alpha.numpy().item(), step=step)
-                         tf.summary.scalar('buffer/beta', replay_buffer.beta, step=step)
-                         tf.summary.scalar('buffer/per_alpha', current_alpha, step=step)
-                         # Log TD error percentiles calculated in Revision 6
-                         if 'step_info_for_logging' in locals():
-                             for k, v in step_info_for_logging.items():
-                                  tf.summary.scalar(f'td_error/{k}', v, step=step)
-                         logger.debug(f"Step {step}: Logged losses, alpha, PER params, TD error stats.")
-                
-            if step % checkpoint_interval == 0:
-                 # Save checkpoint
-                 chkpt_path = os.path.join(self.sac_run_models_dir, f'sac_agent_step_{step}')
-                 meta_data = {'step': step, 'edge_threshold': edge_threshold_heuristic if 'edge_threshold_heuristic' in locals() else self.config.get('calibration',{}).get('edge_threshold')}
-                 agent.save(chkpt_path, meta_data=meta_data)
-                 # --- Save State Filter (Task 5.2) --- #
-                 if self.state_filter:
-                      filter_state = self.state_filter.get_state()
-                      np.savez(os.path.join(chkpt_path, 'state_filter.npz'), **filter_state)
-                 # --- End Save State Filter --- #
-                 logger.info(f"Saved SAC checkpoint at step {step} to {chkpt_path}")
-
-            # Episode end handling
             if done:
-                logger.debug(f"Episode {episode_count} finished after {current_episode_steps} steps. Reward: {current_episode_reward:.2f}. Buffer size: {len(replay_buffer)}.")
-                episode_rewards_log.append({'episode': episode_count, 'total_step': step, 'episode_reward': current_episode_reward, 'episode_steps': current_episode_steps})
-                with summary_writer.as_default(): # Use context manager
-                    tf.summary.scalar('reward/episode', current_episode_reward, step=step)
-                    tf.summary.scalar('steps/episode', current_episode_steps, step=step)
-
-                # Reset for next episode
                 state = env.reset()
-                current_episode_reward = 0.0
+                # Normalize reset state
-                current_episode_steps = 0
-            else:
-                state = next_state
-                
-            # Update progress bar description
-            if step % 100 == 0 and len(replay_buffer) >= min_buffer_size:
-                 last_reward = episode_rewards_log[-1]['episode_reward'] if episode_rewards_log else np.nan
-                 pbar.set_description(f"SAC Training | Ep Reward (Last): {last_reward:.2f} | Buffer: {len(replay_buffer)}")
-                 
-            # Save agent checkpoint
-            if (step + 1) % checkpoint_interval == 0 or step == total_training_steps - 1:
-                 save_path = os.path.join(self.sac_run_models_dir, f'sac_agent_step_{step + 1}')
-                 agent.save_weights(save_path)
-                 logger.info(f"SAC agent weights saved at step {step + 1} to {save_path}")
-                 # --- Save State Filter (Task 5.2) --- #
                 if self.state_filter:
-                      state_filter_path = os.path.join(self.sac_run_models_dir, f'state_filter_step_{step + 1}.npz')
+                     state = self.state_filter(state, update=True)
+                total_reward = 0.0 # Reset episodic reward
+
+            # Save agent checkpoint periodically
+            if step % save_freq == 0 and step > 0:
+                save_path = os.path.join(self.sac_run_models_dir, f'sac_agent_step_{step}')
                 try:
-                           self.state_filter.save_npz(state_filter_path)
+                    agent.save(save_path)
-                           logger.info(f"State filter saved to {state_filter_path}")
+                    # Also save state filter if used
+                    if self.state_filter:
+                         filter_path = os.path.join(save_path, 'state_filter.pkl')
+                         joblib.dump(self.state_filter, filter_path)
+                         logger.info(f"Saved state filter to {filter_path}")
                 except Exception as e:
-                           logger.error(f"Failed to save state filter: {e}")
+                     logger.error(f"Failed to save agent/filter checkpoint at step {step} to {save_path}: {e}", exc_info=True)
-                 # --- End Save State Filter --- #
-                 last_saved_path = save_path # Update last saved path
-                 # Also save the reward log periodically
-                 if episode_rewards_log:
-                      rewards_df = pd.DataFrame(episode_rewards_log)
-                      rewards_log_path = os.path.join(self.sac_run_logs_dir, 'episode_rewards.csv')
-                      try:
-                           rewards_df.to_csv(rewards_log_path, index=False)
-                      except Exception as e:
-                           logger.error(f"Failed to save episode rewards log: {e}")
 
         # --- Final Save --- #
-        pbar.close()
-        summary_writer.close()
         final_save_path = os.path.join(self.sac_run_models_dir, 'sac_agent_final')
-        agent.save_weights(final_save_path)
+        try:
-        logger.info(f"Final SAC agent weights saved to {final_save_path}")
+            agent.save(final_save_path)
-        # Save final state filter
             if self.state_filter:
-             final_state_filter_path = os.path.join(self.sac_run_models_dir, 'state_filter_final.npz')
+                 filter_path = os.path.join(final_save_path, 'state_filter.pkl')
-             try:
+                 joblib.dump(self.state_filter, filter_path)
-                  self.state_filter.save_npz(final_state_filter_path)
+                 logger.info(f"Saved final state filter to {filter_path}")
-                  logger.info(f"Final state filter saved to {final_state_filter_path}")
+            self.last_saved_agent_path = final_save_path # Store path for potential return
         except Exception as e:
-                  logger.error(f"Failed to save final state filter: {e}")
+            logger.error(f"Failed to save final agent/filter checkpoint to {final_save_path}: {e}", exc_info=True)
+            self.last_saved_agent_path = None
 
-        # Save final rewards log
+        end_time = time.time()
-        if episode_rewards_log:
+        training_duration = end_time - start_time
-             rewards_df = pd.DataFrame(episode_rewards_log)
+        logger.info(f"SAC training loop finished in {training_duration:.2f} seconds.")
-             rewards_log_path = os.path.join(self.sac_run_logs_dir, 'episode_rewards.csv')
+        logger.info(f"Final agent checkpoint saved to: {self.last_saved_agent_path}")
+
+        # --- Close imputed transition log file --- #
+        if imputed_log_file:
             try:
-                  rewards_df.to_csv(rewards_log_path, index=False)
+                imputed_log_file.close()
-                  logger.info(f"Final episode rewards log saved to {rewards_log_path}")
+                logger.info(f"Closed imputed transition log file: {imputed_log_path}")
             except Exception as e:
-                  logger.error(f"Failed to save final episode rewards log: {e}")
+                logger.error(f"Error closing imputed transition log file: {e}")
+        # --- End close log file --- #
 
-        return final_save_path if os.path.exists(final_save_path) else last_saved_path
+        return self.last_saved_agent_path
 
     def train(self, gru_run_id_for_sac: str) -> str | None:
         """
@@ -918,12 +1009,10 @@ class SACTrainer:
             initial_lr=self.sac_cfg.get('actor_lr', 3e-4),
             lr_decay_rate=self.sac_cfg.get('lr_decay_rate', 0.96),
             decay_steps=self.sac_cfg.get('decay_steps', 100000),
-            buffer_capacity=self.sac_cfg.get('buffer_max_size', 100000),
             ou_noise_stddev=self.sac_cfg.get('ou_noise_stddev', 0.2),
             alpha=self.sac_cfg.get('alpha', 0.2),
             alpha_auto_tune=self.sac_cfg.get('alpha_auto_tune', True),
             target_entropy=self.sac_cfg.get('target_entropy', -1.0 * env.action_dim),
-            min_buffer_size=self.sac_cfg.get('min_buffer_size', 1000),
             edge_threshold_config=current_edge_threshold, # Pass edge threshold
             # --- Pass Env Params (Task 5.6) --- #
             reward_scale_config=reward_scale,

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), \
+        assert len(mu_predictions) == len(sigma_predictions) == len(p_cal_predictions) == len(actual_returns) == len(bar_imputed_flags), \
-            "All input arrays must have the same length"
+            "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]
 
+        # 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
-        z_score_t = np.abs(mu_t) / (sigma_t + 1e-9)
+        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