"""
SAC Trainer Component.

Handles the offline training process for the SAC agent, including loading 
dependencies from a specified GRU run, preparing data for the environment, 
and executing the training loop.
"""

import tensorflow as tf
import json
import logging
import os
import sys
import yaml
import joblib
import pandas as pd
import numpy as np
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
from gru_sac_predictor.src.data_loader import DataLoader
from gru_sac_predictor.src.feature_engineer import FeatureEngineer
from gru_sac_predictor.src.gru_model_handler import GRUModelHandler
from gru_sac_predictor.src.calibrator import Calibrator
from gru_sac_predictor.src.sac_agent import SACTradingAgent
from gru_sac_predictor.src.trading_env import TradingEnv
try:
    from gru_sac_predictor.src.features import minimal_whitelist # For FE fallback
except ImportError:
    minimal_whitelist = [] # Define empty if import fails

# --- Import MeanStdFilter (Task 5.2) --- #
try:
    from gru_sac_predictor.src.utils.running_stats import MeanStdFilter
    STATE_FILTER_AVAILABLE = True
except ImportError:
     logging.warning("MeanStdFilter not found in utils. State normalization will be disabled.")
     MeanStdFilter = None
     STATE_FILTER_AVAILABLE = False
# --- End Import --- #

logger = logging.getLogger(__name__)

# ==============================================================================
# Prioritized Experience Replay (PER) Buffer Implementation
# ==============================================================================

class SumTree:
    """ Simple SumTree implementation for PER priority sampling."""
    write = 0

    def __init__(self, capacity):
        self.capacity = capacity
        # Tree structure: Internal nodes + leaves
        self.tree = np.zeros(2 * capacity - 1)
        # Data storage (corresponds to leaves)
        self.data = np.zeros(capacity, dtype=object) 
        self.n_entries = 0

    def _propagate(self, idx, change):
        parent = (idx - 1) // 2
        self.tree[parent] += change
        if parent != 0:
            self._propagate(parent, change)

    def _retrieve(self, idx, s):
        left = 2 * idx + 1
        right = left + 1

        if left >= len(self.tree):
            return idx

        if s <= self.tree[left]:
            return self._retrieve(left, s)
        else:
            return self._retrieve(right, s - self.tree[left])

    def total(self):
        return self.tree[0]

    def add(self, p, data):
        # Data is now expected to be (sample, is_seeded)
        idx = self.write + self.capacity - 1

        self.data[self.write] = data
        self.update(idx, p)

        self.write += 1
        if self.write >= self.capacity:
            self.write = 0

        if self.n_entries < self.capacity:
            self.n_entries += 1

    def update(self, idx, p):
        change = p - self.tree[idx]
        self.tree[idx] = p
        self._propagate(idx, change)

    def get(self, s):
        idx = self._retrieve(0, s)
        dataIdx = idx - self.capacity + 1
        # Return index, priority, and the stored (sample, is_seeded) tuple
        return (idx, self.tree[idx], self.data[dataIdx])


class PrioritizedReplayBuffer:
    """ PER buffer using a SumTree."""
    epsilon = 0.01  # Small amount to avoid zero priority
    beta = 0.4  # Importance sampling exponent
    beta_increment_per_sampling = 0.001
    abs_err_upper = 1.  # Clipped abs error

    def __init__(self, capacity, alpha=0.6, beta_start=0.4, beta_frames=100000):
        self.tree = SumTree(capacity)
        self.capacity = capacity
        # Alpha is annealed externally now
        self.beta = beta_start
        # Calculate beta increment based on total frames/steps
        self.beta_increment_per_sampling = (1. - beta_start) / beta_frames 
        self.max_priority = 1.0 # Initialize max priority

        # --- Revision 5: Oracle Seed Decay Steps --- #
        self.per_seed_decay_steps = beta_frames // 2 # Default: decay over half training
        logger.info(f"PER Oracle Seed IS weight decay steps: {self.per_seed_decay_steps}")
        # --- End Revision 5 --- #

    def add(self, error, sample, is_seeded=False):
        """ Store new experience with initial priority and seeded flag. """
        p = self.max_priority # Use max priority for new samples
        # Store sample along with its seeded status
        self.tree.add(p, (sample, is_seeded))

    def sample(self, n, beta=None):
        """ Sample batch, returning indices, samples, and IS weights. """
        if beta is None: 
             beta = self.beta # Use current beta if not specified
             
        batch = []
        idxs = []
        segment = self.tree.total() / n
        priorities = []

        # Anneal beta
        self.beta = np.min([1., self.beta + self.beta_increment_per_sampling])

        for i in range(n):
            a = segment * i
            b = segment * (i + 1)
            s = np.random.uniform(a, b)
            (idx, p, data) = self.tree.get(s)
            priorities.append(p)
            # Unpack data: contains (sample, is_seeded)
            if isinstance(data, tuple) and len(data) == 2:
                 batch.append(data) # Keep tuple for now
            else: # Fallback for older buffer states maybe?
                 batch.append((data, False)) # Assume not seeded
            idxs.append(idx)

        sampling_probabilities = np.array(priorities) / self.tree.total()
        is_weight = np.power(self.tree.n_entries * sampling_probabilities, -beta)
        is_weight /= is_weight.max() # Normalize weights

        return idxs, batch, is_weight

    # --- Revision 5: Modified priority calculation and update --- #
    def _get_priority(self, error, alpha):
        """Calculate priority: (clip(|error|) + eps) ^ alpha."""
        # Clip absolute error
        clipped_abs_error = np.clip(np.abs(error), 0, self.abs_err_upper)
        return (clipped_abs_error + self.epsilon) ** alpha

    def update_priorities(self, idxs, errors, alpha):
        """ Update priorities of sampled transitions using current alpha. """
        for i, idx in enumerate(idxs):
            # Use current annealed alpha
            p = self._get_priority(errors[i], alpha)
            self.max_priority = max(self.max_priority, p) # Keep track of max priority
            self.tree.update(idx, p)
    # --- End Revision 5 --- #
             
    def __len__(self):
        return self.tree.n_entries

# ==============================================================================

class SACTrainer:
    """Manages the offline SAC training workflow."""

    def __init__(self, config: dict, base_models_dir: str, base_logs_dir: str, base_results_dir: str):
        """
        Initialize the SACTrainer.

        Args:
            config (dict): The main pipeline configuration dictionary.
            base_models_dir (str): The base directory where all run models are stored (e.g., project_root/models).
            base_logs_dir (str): Base directory for logs.
            base_results_dir (str): Base directory for results.
        """
        self.config = config
        self.base_models_dir = base_models_dir
        self.sac_cfg = config['sac']
        self.env_cfg = config.get('environment', {})
        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)
        self.per_beta_start = self.sac_cfg.get('per_beta_start', 0.4)
        # Estimate beta annealing frames based on total steps
        total_training_steps = self.sac_cfg.get('total_training_steps', 100000)
        self.per_beta_frames = self.sac_cfg.get('per_beta_frames', total_training_steps)
        # --- End PER Config --- #

        # Generate a specific run ID for this SAC training instance
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.sac_train_run_id = f"sac_train_{timestamp}"
        logger.info(f"Initializing SACTrainer with Run ID: {self.sac_train_run_id}")

        # Setup directories for this specific SAC training run
        self.sac_run_models_dir = os.path.join(self.base_models_dir, self.sac_train_run_id)
        self.sac_run_logs_dir = os.path.join(base_logs_dir, self.sac_train_run_id)
        self.sac_run_results_dir = os.path.join(base_results_dir, self.sac_train_run_id)
        self.sac_tb_log_dir = os.path.join(self.sac_run_logs_dir, 'tensorboard')

        os.makedirs(self.sac_run_models_dir, exist_ok=True)
        os.makedirs(self.sac_run_logs_dir, exist_ok=True)
        os.makedirs(self.sac_run_results_dir, exist_ok=True)
        os.makedirs(self.sac_tb_log_dir, exist_ok=True)
        
        # --- Initialize State Filter (Task 5.2) --- #
        # Get state dim - assumes TradingEnv has a fixed state_dim attribute
        # We might need to instantiate a dummy env briefly to get this?
        # Or hardcode based on current TradingEnv._get_state()
        # Current state: [mu, sigma, edge, |mu|/sigma, position] -> dim=5
        # TODO: Get state_dim more robustly if env changes
        self.state_dim_env = 5 # Hardcoded based on current TradingEnv
        self.state_filter = None
        if STATE_FILTER_AVAILABLE and self.config.get('sac',{}).get('use_state_filter', True): # Add config flag
             logger.info(f"Initializing MeanStdFilter for state normalization (shape={self.state_dim_env}).")
             self.state_filter = MeanStdFilter(shape=(self.state_dim_env,))
        else:
             logger.warning("State filter is disabled (either unavailable or config flag is false).")
        # --- End Initialize State Filter --- #

        # Configure logging specifically for this trainer instance if needed
        # For now, relies on the pipeline's logger setup
        logger.info(f"  SAC Models Dir: {self.sac_run_models_dir}")
        logger.info(f"  SAC Logs Dir:   {self.sac_run_logs_dir}")
        logger.info(f"  SAC Results Dir:{self.sac_run_results_dir}")
        logger.info(f"  SAC TB Dir:     {self.sac_tb_log_dir}")
        
        # Save config subset relevant to SAC training?
        # Or assume full config is saved by the main pipeline

    def _load_gru_dependencies(self, gru_run_id: str) -> dict | None:
        """
        Loads artifacts (whitelist, scaler, GRU model, T) from a completed GRU pipeline run.

        Args:
            gru_run_id (str): The run ID of the GRU pipeline run.

        Returns:
            dict | None: A dictionary containing the loaded dependencies 
                         ('whitelist', 'scaler', 'gru_model', 'optimal_T'), or None on failure.
        """
        logger.info(f"--- Loading Dependencies from GRU Run ID: {gru_run_id} ---")
        gru_run_models_dir = os.path.join(self.base_models_dir, gru_run_id)
        if not os.path.exists(gru_run_models_dir):
             logger.error(f"Models directory for GRU run {gru_run_id} not found at: {gru_run_models_dir}")
             return None
             
        dependencies = {}

        # 1. Load Whitelist
        whitelist_path = os.path.join(gru_run_models_dir, f"final_whitelist_{gru_run_id}.json")
        try:
            with open(whitelist_path, 'r') as f:
                dependencies['whitelist'] = json.load(f)
            logger.info(f"Loaded whitelist ({len(dependencies['whitelist'])} features) from {whitelist_path}")
            if not dependencies['whitelist']:
                 raise ValueError("Loaded whitelist is empty.")
        except Exception as e:
            logger.error(f"Failed to load whitelist from {whitelist_path}: {e}", exc_info=True)
            return None

        # 2. Load Scaler
        scaler_path = os.path.join(gru_run_models_dir, f"feature_scaler_{gru_run_id}.joblib")
        try:
            dependencies['scaler'] = joblib.load(scaler_path)
            logger.info(f"Loaded scaler from {scaler_path}")
        except Exception as e:
            logger.error(f"Failed to load scaler from {scaler_path}: {e}", exc_info=True)
            return None

        # 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
        # Pass self.config (the trainer's config) to the GRUModelHandler
        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"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 Calibration Info and Parameters
        calib_info_path = os.path.join(gru_run_models_dir, f"calibration_info_{gru_run_id}.json")
        calib_params = None
        calib_method = None
        try:
            with open(calib_info_path, 'r') as f:
                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 calibration info/parameters: {e}", exc_info=True)
            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

    def _prepare_data_for_sac(self, gru_dependencies: dict) -> tuple | None:
        """
        Replicates the necessary data loading and preparation steps using the loaded 
        GRU dependencies, specifically focusing on the VALIDATION dataset to 
        generate inputs for the TradingEnv.

        Args:
            gru_dependencies (dict): The dictionary returned by _load_gru_dependencies.

        Returns:
            tuple | None: A tuple containing (mu_val, sigma_val, p_cal_val, actual_ret_val)
                          for the validation set, or None on failure.
        """
        logger.info("--- Preparing Validation Data for SAC Environment --- ")
        try:
            # 1. Load Raw Data (using a temporary DataLoader)
            temp_data_loader = DataLoader(db_dir=self.data_cfg['db_dir'])
            df_raw = temp_data_loader.load_data(
                ticker=self.data_cfg['ticker'],
                exchange=self.data_cfg['exchange'],
                start_date=self.data_cfg['start_date'],
                end_date=self.data_cfg['end_date'],
                interval=self.data_cfg['interval']
            )
            if df_raw is None or df_raw.empty: raise ValueError("Raw data loading failed")
            df_raw.dropna(subset=['open', 'high', 'low', 'close', 'volume'], inplace=True)
            logger.info("Loaded raw data.")

            # 2. Engineer Base Features (using config)
            # FIX: Instantiate FeatureEngineer correctly using self.config
            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. Define Labels (on the full engineered data)
            horizon = self.config['gru'].get('prediction_horizon', 5)
            target_ret_col = f'fwd_log_ret_{horizon}'
            target_dir_col = f'direction_label3_{horizon}' if self.use_ternary else f'direction_label_{horizon}'
            _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)

            # 4. Align Features and Targets
            df_engineered.dropna(subset=[target_ret_col, target_dir_col], inplace=True)
            # Identify *all* potential feature columns (excluding targets)
            potential_feature_cols = [col for col in df_engineered.columns if col not in [target_ret_col, target_dir_col, 'future_close']]
            # Ensure whitelist features are present in potential features
            loaded_whitelist = gru_dependencies['whitelist']
            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 (using aligned engineered data)
            wf_enabled = self.config.get('walk_forward', {}).get('enabled', False)
            if wf_enabled: logger.warning("Walk-forward enabled, but SAC uses split_ratios.")
            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_engineered.index[train_end_idx:val_end_idx]
            if val_indices.empty: raise ValueError("Validation split resulted in empty indices.")
            df_val_aligned = df_engineered.loc[val_indices]

            # -- Determine columns expected by scaler --
            scaler = gru_dependencies['scaler']
            expected_scaler_features = []
            if hasattr(scaler, 'feature_names_in_'):
                expected_scaler_features = scaler.feature_names_in_.tolist()
                logger.debug(f"Scaler expects features: {expected_scaler_features}")
            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}")

            # 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_list, y_val_seq_targets_list, val_seq_indices = [], [], []
            features_np_arr = X_val_pruned_scaled.values
            targets_np_arr = y_val.values
            for i in range(lookback, len(features_np_arr)):
                X_val_seq_list.append(features_np_arr[i-lookback : i])
                y_val_seq_targets_list.append(targets_np_arr[i])
                val_seq_indices.append(y_val.index[i])

            if not X_val_seq_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)
            actual_ret_val_seq = y_val_seq_targets[:, 0]
            y_dir_val_seq = y_val_seq_targets[:, 1]
            logger.info(f"Created validation sequences (X shape: {X_val_seq.shape}).")

            # 10. Get GRU Predictions using loaded GRU model directly
            gru_model = gru_dependencies['gru_model']
            logger.info(f"Generating GRU predictions using loaded model (type: {type(gru_model)}).")

            # Check model type and predict accordingly
            if not hasattr(gru_model, 'predict'):
                 raise TypeError("Loaded GRU model object does not have a 'predict' method.")

            predictions_val = gru_model.predict(X_val_seq)

            # --- Explicit Shape Logging --- #
            if isinstance(predictions_val, list):
                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 \
                    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}, "
                                  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)}")

            # 11. Calibrate Predictions using loaded parameters
            calib_method = gru_dependencies.get('calibration_method')
            calib_params = gru_dependencies.get('calibration_params')
            p_cal_val_pred = None

            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

        except Exception as e:
            logger.error(f"Error preparing data for SAC environment: {e}", exc_info=True)
            return None

    def _load_agent_for_resume(self, agent: SACTradingAgent) -> None:
        """Loads agent weights and state filter status if resuming."""
        load_run_id = self.control_cfg.get('sac_resume_run_id')
        load_step = self.control_cfg.get('sac_resume_step', 'final')
        current_edge_threshold = self.config.get('calibration', {}).get('edge_threshold', 0.55)

        if not load_run_id:
            logger.info("No SAC resume run ID specified. Starting training from scratch.")
            return

        # Construct path relative to base models dir
        if load_step == 'final':
            load_path = os.path.join(self.base_models_dir, f"{load_run_id}", 'sac_agent_final')
        else:
            try:
                step_num = int(load_step)
                load_path = os.path.join(self.base_models_dir, f"{load_run_id}", f'sac_agent_step_{step_num}')
            except ValueError:
                logger.error(f"Invalid sac_resume_step: {load_step}. Must be 'final' or an integer. Starting fresh.")
                return

        logger.info(f"Attempting to load SAC agent from {load_path} to resume training...")
        if os.path.exists(load_path):
            try:
                loaded_meta = agent.load(load_path)
                # Check for Buffer Purge on Load 
                saved_edge_thr = loaded_meta.get('edge_threshold')
                if saved_edge_thr is not None and abs(saved_edge_thr - current_edge_threshold) > 1e-6:
                    logger.warning(f'Edge threshold mismatch on load (Saved={saved_edge_thr:.3f}, Current={current_edge_threshold:.3f}). Clearing replay buffer before resuming.')
                    agent.clear_buffer()
                elif saved_edge_thr is None:
                    logger.warning("Loaded SAC agent metadata did not contain 'edge_threshold'. Cannot verify consistency.")
                else:
                    logger.info('Edge threshold consistent with loaded agent metadata.')
                    
                # --- Load State Filter (Task 5.2) --- #
                filter_state_path = os.path.join(load_path, 'state_filter.npz')
                if self.state_filter is not None and os.path.exists(filter_state_path):
                     try:
                          with np.load(filter_state_path) as data:
                               filter_state = {key: data[key] for key in data.files}
                          self.state_filter.set_state(filter_state)
                          logger.info(f"Loaded state filter state from {filter_state_path}")
                     except Exception as filter_e:
                          logger.error(f"Failed to load state filter state from {filter_state_path}: {filter_e}. Filter will be reset.")
                elif self.state_filter is not None:
                     logger.warning(f"State filter state file not found at {filter_state_path}. Filter will be reset.")
                # --- End Load State Filter --- #

            except Exception as e:
                logger.error(f"Failed to load SAC agent for resume: {e}. Starting fresh.", exc_info=True)
        else:
            logger.warning(f"SAC agent path not found for resume: {load_path}. Starting fresh.")
            
    def _training_loop(self, agent: SACTradingAgent, env: TradingEnv) -> str | None:
        """The main SAC training loop."""
        total_steps = self.sac_cfg.get('total_training_steps', 100000)
        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)

        # Initialize Replay Buffer (Standard or Prioritized)
        if self.use_per:
            logger.info(f"Using Prioritized Replay Buffer (Capacity: {buffer_capacity})")
            replay_buffer = PrioritizedReplayBuffer(
                capacity=buffer_capacity, 
                alpha=self.per_alpha, # Initial alpha 
                beta_start=self.per_beta_start,
                beta_frames=self.per_beta_frames
            )
        else:
            logger.info(f"Using Standard Replay Buffer (Capacity: {buffer_capacity})")
            replay_buffer = collections.deque(maxlen=buffer_capacity)

        # TensorBoard setup
        tb_callback = TensorBoard(log_dir=self.sac_tb_log_dir)
        # --- Revision 4: Set model for TensorBoard --- #
        # Check if agent has actor/critic models accessible
        # This depends heavily on SACTradingAgent implementation
        # Assuming agent.actor and agent.critic1/2 are the models
        if hasattr(agent, 'actor') and hasattr(agent, 'critic1') and hasattr(agent, 'critic2'):
             tb_callback.set_model(agent.actor) # Link to one model is often sufficient 
             logger.info("TensorBoard callback linked to SAC agent model (actor).")
        else:
             logger.warning("Could not link TensorBoard callback to agent models (actor/critic not found).")
        # --- 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:
            imputed_log_file = open(imputed_log_path, 'w', newline='')
            imputed_csv_writer = csv.writer(imputed_log_file)
            imputed_csv_writer.writerow(['step', 'imputed_handling_mode', 'action', 'reward', 'position_before', 'position_after'])
            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 --- #

        state = env.reset()
        # Normalize initial state if filter is active
        if self.state_filter:
            state = self.state_filter(state, update=True) # Update filter with initial state

        total_reward = 0.0
        start_time = time.time()

        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:
                action = agent.select_action(state)

            # Step the environment
            next_state_raw, reward, done, info = env.step(action)

            # Check if the step was skipped due to imputed bar handling
            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:
                    # Add with initial error=1 (or max priority) and is_seeded=False
                    # The error will be updated after the first training step on this sample.
                    replay_buffer.add(error=1.0, sample=(state_before_action, action, reward, next_state, done))
                else:
                    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
            total_reward += reward

            # Perform SAC agent updates
            if step >= update_after and step % update_every == 0:
                for update_i in range(update_every): # Perform multiple updates per interval
                    if self.use_per:
                        if len(replay_buffer) > batch_size:
                            idxs, batch_data, is_weights = replay_buffer.sample(batch_size, beta=replay_buffer.beta) # Use annealed beta
                            # Unpack batch_data which contains (sample, is_seeded) tuples
                            batch = [item[0] for item in batch_data]
                            update_info = agent.update(batch, is_weights=is_weights, per_beta=replay_buffer.beta)
                            if update_info and 'td_errors' in update_info:
                                 # Update priorities using current annealed alpha
                                 current_alpha = agent.get_current_per_alpha(step)
                                 replay_buffer.update_priorities(idxs, update_info['td_errors'], alpha=current_alpha)
                        else:
                            continue # Not enough samples yet for PER
                    else: # Standard buffer
                        if len(replay_buffer) > batch_size:
                            indices = np.random.choice(len(replay_buffer), size=batch_size, replace=False)
                            batch = [replay_buffer[i] for i in indices]
                            update_info = agent.update(batch)
                        else:
                             continue # Not enough samples yet

                    # Log training metrics (losses, Q-values, alpha) to TensorBoard
                    if update_info and step % log_freq == 0 and update_i == 0: # Log once per interval
                        with tb_callback.writer.as_default():
                            for key, value in update_info.items():
                                if key != 'td_errors': # Don't log TD errors directly
                                     tf.summary.scalar(f'sac/{key}', value, step=step)
                        # logger.debug(f"Step {step}: Logged SAC metrics to TensorBoard.")

            # Check environment done state
            if done:
                state = env.reset()
                # Normalize reset state
                if self.state_filter:
                     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:
                    agent.save(save_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 agent/filter checkpoint at step {step} to {save_path}: {e}", exc_info=True)

        # --- Final Save --- #
        final_save_path = os.path.join(self.sac_run_models_dir, 'sac_agent_final')
        try:
            agent.save(final_save_path)
            if self.state_filter:
                 filter_path = os.path.join(final_save_path, 'state_filter.pkl')
                 joblib.dump(self.state_filter, filter_path)
                 logger.info(f"Saved final state filter to {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 agent/filter checkpoint to {final_save_path}: {e}", exc_info=True)
            self.last_saved_agent_path = None

        end_time = time.time()
        training_duration = end_time - start_time
        logger.info(f"SAC training loop finished in {training_duration:.2f} seconds.")
        logger.info(f"Final agent checkpoint saved to: {self.last_saved_agent_path}")

        # --- Close imputed transition log file --- #
        if imputed_log_file:
            try:
                imputed_log_file.close()
                logger.info(f"Closed imputed transition log file: {imputed_log_path}")
            except Exception as e:
                logger.error(f"Error closing imputed transition log file: {e}")
        # --- End close log file --- #

        return self.last_saved_agent_path

    def train(self, gru_run_id_for_sac: str) -> str | None:
        """
        Main entry point to start the SAC training process.

        Args:
            gru_run_id_for_sac (str): The run ID of the GRU pipeline run whose artifacts should be used.
        
        Returns:
            str | None: Path to the final saved SAC agent, or None if training failed.
        """
        logger.info(f"=== Starting SAC Training Process (SAC Run ID: {self.sac_train_run_id}) ===")
        logger.info(f"Using artifacts from GRU Run ID: {gru_run_id_for_sac}")

        # 1. Load GRU dependencies
        gru_dependencies = self._load_gru_dependencies(gru_run_id_for_sac)
        if gru_dependencies is None:
             logger.error("Failed to load GRU dependencies. Aborting SAC training.")
             return None

        # 2. Prepare data for SAC environment (using validation set)
        env_data = self._prepare_data_for_sac(gru_dependencies)
        if env_data is None:
             logger.error("Failed to prepare data for SAC environment. Aborting SAC training.")
             return None
        mu_val, sigma_val, p_cal_val, actual_ret_val = env_data

        # 3. Initialize Environment
        logger.info("Initializing Trading Environment...")
        env = TradingEnv(
            mu_predictions=mu_val,
            sigma_predictions=sigma_val,
            p_cal_predictions=p_cal_val,
            actual_returns=actual_ret_val,
            initial_capital=self.env_cfg.get('initial_capital', 10000.0),
            transaction_cost=self.env_cfg.get('transaction_cost', 0.0005)
        )
        logger.info(f"TradingEnv initialized with {env.n_steps} steps.")

        # 4. Initialize SAC Agent
        logger.info("Initializing SAC Agent...")
        current_edge_threshold = self.config.get('calibration', {}).get('edge_threshold', 0.55)
        # --- Get Env Params for Agent Metadata (Task 5.6) --- #
        reward_scale = self.env_cfg.get('reward_scale', 100.0) # Default from TradingEnv
        action_penalty_lambda = self.env_cfg.get('action_penalty_lambda', 0.0) # Default from TradingEnv
        # --- End Get Env Params --- #
        agent = SACTradingAgent(
            state_dim=env.state_dim,
            action_dim=env.action_dim,
            gamma=self.sac_cfg.get('gamma', 0.99),
            tau=self.sac_cfg.get('tau', 0.005),
            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),
            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),
            edge_threshold_config=current_edge_threshold, # Pass edge threshold
            # --- Pass Env Params (Task 5.6) --- #
            reward_scale_config=reward_scale,
            action_penalty_lambda_config=action_penalty_lambda
            # --- End Pass Env Params --- #
        )
        logger.info("SAC Agent initialized.")

        # 5. Load agent weights if resuming
        self._load_agent_for_resume(agent)

        # --- Add Training Size Assertion (Step 4-C) --- #
        total_steps_cfg = self.sac_cfg.get('total_training_steps', 100000)
        val_len = len(actual_ret_val) # Length of the validation dataset used by env
        min_required_steps = int(0.4 * val_len)
        assert total_steps_cfg >= min_required_steps, \
            f"Configured total_training_steps ({total_steps_cfg}) is less than 40% of the validation set size ({val_len}). Minimum required: {min_required_steps}. Consider increasing training steps."
        logger.info(f"Training size check passed: {total_steps_cfg=}, {min_required_steps=}")
        # --- End Assertion --- #

        # 6. Run training loop
        final_agent_path = self._training_loop(agent, env)

        if final_agent_path:
            logger.info(f"=== SAC Training Process Completed Successfully ===")
        else:
            logger.error("=== SAC Training Process Failed ===")
            
        return final_agent_path 