import numpy as np
import pandas as pd
import os
from tqdm import tqdm # Added for progress tracking
import logging # Added for feature functions
from scipy import stats # Added for feature functions
from typing import Tuple
import tensorflow as tf
import warnings
import sys

# Optional: Import TA-Lib if available (used by V6 features)
try:
    import talib
    TALIB_AVAILABLE = True
    # logging.info("TA-Lib found. Using TA-Lib for V6 features.") # Reduce noise
except ImportError:
    TALIB_AVAILABLE = False
    logging.warning("TA-Lib not found. V6 feature calculation will be limited.")

# Import other components from the src directory
from .gru_predictor import CryptoGRUModel
from .sac_agent_simplified import SimplifiedSACTradingAgent
# Remove direct import of prepare_features as logic is integrated here
# from .data_pipeline import prepare_features

# Optional: Matplotlib for plotting
try:
    import matplotlib.pyplot as plt
except ImportError:
    plt = None

# --- V7.12 START: Add SAC Training History Plotting Function ---
def plot_sac_training_history(history, save_path=None, run_id=None):
    """
    Plots the actor and critic loss from the SAC training history.

    Args:
        history (list[dict]): A list of dictionaries, where each dict contains
                              metrics like 'actor_loss' and 'critic_loss'.
        save_path (str, optional): Path to save the plot. If None, tries to show plot.
        run_id (str, optional): Run ID to include in the plot title.
    """
    if plt is None:
        print("Matplotlib not available, skipping SAC training plot.")
        return

    if not history:
        print("No SAC training history data provided, skipping plot.")
        return

    try:
        actor_losses = [item['actor_loss'] for item in history if 'actor_loss' in item and item['actor_loss'] is not None]
        critic_losses = [item['critic_loss'] for item in history if 'critic_loss' in item and item['critic_loss'] is not None]
        steps = range(1, len(history) + 1) # Assuming one entry per training step/batch

        if not actor_losses or not critic_losses:
            print("Warning: Could not extract valid actor or critic losses from history. Skipping plot.")
            return

        # Determine common length if lists differ (shouldn't usually happen if logged together)
        min_len = min(len(actor_losses), len(critic_losses))
        if min_len < len(history):
             print(f"Warning: Plotting only {min_len} steps due to missing data.")
             steps = range(1, min_len + 1)
             actor_losses = actor_losses[:min_len]
             critic_losses = critic_losses[:min_len]

        fig, axs = plt.subplots(2, 1, figsize=(12, 10), sharex=True)
        fig.suptitle(f'SAC Training History {"(Run: " + run_id + ")" if run_id else ""}', fontsize=16)

        # Actor Loss Plot
        axs[0].plot(steps, actor_losses, label='Actor Loss', color='tab:blue')
        axs[0].set_ylabel('Loss')
        axs[0].set_title('Actor Loss over Training Steps')
        axs[0].legend()
        axs[0].grid(True)

        # Critic Loss Plot
        axs[1].plot(steps, critic_losses, label='Critic Loss', color='tab:orange')
        axs[1].set_xlabel('Agent Training Step')
        axs[1].set_ylabel('Loss')
        axs[1].set_title('Critic Loss over Training Steps')
        axs[1].legend()
        axs[1].grid(True)

        plt.tight_layout(rect=[0, 0.03, 1, 0.95]) # Adjust layout to prevent title overlap

        if save_path:
            try:
                plt.savefig(save_path)
                print(f"SAC training history plot saved to {save_path}")
            except Exception as e:
                print(f"Error saving SAC training plot: {e}")
        else:
            try:
                print("Displaying SAC training plot...")
                plt.show()
            except Exception as e:
                 print(f"Error displaying SAC training plot: {e}")

    except KeyError as e:
         print(f"KeyError accessing loss data in history: {e}. Skipping plot.")
    except Exception as e:
         print(f"An unexpected error occurred during SAC plot generation: {e}")
    finally:
         # Ensure plot is closed to free memory, especially if not shown interactively
         if plt:
             plt.close(fig)
# --- V7.12 END: Add SAC Training History Plotting Function ---


# --- V6 Feature Calculation Logic (Copied & Adapted) ---
# Configure logging for feature calculation functions
feature_logger = logging.getLogger('TradingSystemFeatures')
# Set level to WARNING to avoid excessive logs during backtest
feature_logger.setLevel(logging.WARNING) 

# V7 Update: Add Scalers and sequence creation
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from .data_pipeline import create_sequences_v2

# Copied from v6/src/cryptofeatures.py
def calculate_vwap(df: pd.DataFrame, period: int = None) -> pd.Series:
    """
    Calculate Volume Weighted Average Price (VWAP).
    Adapted from V6.
    """
    try:
        typical_price = (df['high'] + df['low'] + df['close']) / 3
        tp_volume = typical_price * df['volume']
        volume_sum = df['volume'].rolling(window=period, min_periods=1).sum()
        tp_volume_sum = tp_volume.rolling(window=period, min_periods=1).sum()
        volume_sum_safe = volume_sum.replace(0, np.nan)
        vwap = tp_volume_sum / volume_sum_safe
        vwap = vwap.fillna(method='ffill').fillna(method='bfill').fillna(df['close']) # More robust fill
        return vwap
    except Exception as e:
        feature_logger.error(f"Failed to calculate VWAP (period={period}): {e}")
        return df['close'] # Fallback

# Copied from v6/src/cryptofeatures.py and adapted
def add_crypto_features(df: pd.DataFrame) -> pd.DataFrame:
    """Add cryptocurrency-specific features. Adapted from V6."""
    df_features = df.copy()
    try:
        # Parkinson Volatility - Add better safeguards
        high_safe = df_features['high'].replace(0, np.nan)
        low_safe = df_features['low'].replace(0, np.nan).fillna(high_safe * 0.999)  # Avoid division by zero
        high_to_low = np.clip(high_safe / low_safe, 1.0, 10.0)  # Clip extreme ratios
        log_hl_sq = np.log(high_to_low)**2
        df_features['parkinson_vol_14'] = np.sqrt((1 / (4 * np.log(2))) * log_hl_sq.rolling(window=14).mean())
        
        # Garman-Klass Volatility - Add better safeguards
        close_safe = df_features['close'].replace(0, np.nan)
        open_safe = df_features['open'].replace(0, np.nan).fillna(close_safe * 0.999)  # Avoid division by zero
        close_to_open = np.clip(close_safe / open_safe, 0.5, 2.0)  # Clip extreme ratios
        log_co_sq = np.log(close_to_open)**2
        gk_vol = (0.5 * log_hl_sq - (2 * np.log(2) - 1) * log_co_sq).rolling(window=14).mean()
        # Ensure non-negative values before sqrt
        gk_vol = np.maximum(gk_vol, 0)
        df_features['garman_klass_vol_14'] = np.sqrt(gk_vol)

        # VWAP Features - Improve robustness
        for period in [30, 60, 120]:
            df_features[f'vwap_{period}'] = calculate_vwap(df_features, period=period)
            vwap_safe = df_features[f'vwap_{period}'].replace(0, np.nan).fillna(close_safe)
            # Clip the ratio to a reasonable range to avoid extreme values
            df_features[f'close_to_vwap_{period}'] = np.clip((close_safe / vwap_safe) - 1, -0.5, 0.5)

        # V7 Fix: Cyclical features are now calculated earlier in main.py
        # # Cyclical Features (Simplified)
        # if isinstance(df_features.index, pd.DatetimeIndex):
        #     df_features['hour_sin'] = np.sin(2 * np.pi * df_features.index.hour / 24)
        #     df_features['hour_cos'] = np.cos(2 * np.pi * df_features.index.hour / 24)

        # Volume Intensity - Improve stability
        vol = df_features['volume'].replace(0, np.nan)
        vol_mean_30 = vol.rolling(window=30).mean()
        vol_mean_30_safe = vol_mean_30.replace(0, np.nan).fillna(vol.mean())
        df_features['vol_intensity'] = np.clip(vol / vol_mean_30_safe, 0, 10)  # Clip to reasonable range

        # Price Pattern Features - Improved stability
        body = abs(close_safe - open_safe)
        # Use a meaningful minimum body size to avoid extreme ratios
        min_body_size = df_features['close'].mean() * 0.0001  # Small percentage of avg price
        body_safe = np.maximum(body, min_body_size)
        
        # Calculate wicks with better handling of extreme cases
        upper_wick = df_features['high'] - df_features[['open', 'close']].max(axis=1)
        lower_wick = df_features[['open', 'close']].min(axis=1) - df_features['low']
        # Clip the ratios to reasonable ranges
        df_features['upper_wick_ratio'] = np.clip(upper_wick / body_safe, 0, 5)
        df_features['lower_wick_ratio'] = np.clip(lower_wick / body_safe, 0, 5)

        # Fill NaNs introduced here
        cols_to_fill = ['parkinson_vol_14', 'garman_klass_vol_14', 'vol_intensity', 
                        'upper_wick_ratio', 'lower_wick_ratio'] # Removed hour_sin/cos
        cols_to_fill.extend([f'close_to_vwap_{p}' for p in [30, 60, 120]])
        for col in cols_to_fill:
            if col in df_features.columns:
                 # Fill NaNs with median values rather than zeros to maintain distribution
                 median_val = df_features[col].median()
                 df_features[col].fillna(median_val if not np.isnan(median_val) else 0, inplace=True)

    except Exception as e:
        feature_logger.error(f"Error calculating crypto-specific features: {e}", exc_info=True)
    return df_features

# Adapted from v6/src/data_preprocessing.py::calculate_technical_indicators
def calculate_v6_features(df: pd.DataFrame) -> pd.DataFrame:
    """Calculates V6 technical indicators + basic return features."""
    df_features = df.copy()
    required_cols = ['open', 'high', 'low', 'close', 'volume']
    if not all(col in df_features.columns for col in required_cols):
        missing = [col for col in required_cols if col not in df_features.columns]
        feature_logger.error(f"Missing required V6 columns: {missing}")
        return df_features
    for col in required_cols:
        df_features[col] = pd.to_numeric(df_features[col], errors='coerce')
        if df_features[col].isnull().any():
            df_features[col] = df_features[col].ffill().bfill()
            if df_features[col].isnull().any(): df_features[col] = df_features[col].fillna(0)

    # --- V7.2 Add Past Return Features ---
    for lag in [1, 5, 15, 60]: # 1m, 5m, 15m, 1h returns
        # Use pct_change for robustness to price levels
        df_features[f'return_{lag}m'] = df_features['close'].pct_change(periods=lag)
    # --- End Return Features ---

    if TALIB_AVAILABLE:
        try:
            close = df_features['close'].values; open_price = df_features['open'].values
            high = df_features['high'].values; low = df_features['low'].values; volume = df_features['volume'].values
            for period in [5, 10, 20, 30, 50, 100]:
                df_features[f'SMA_{period}'] = talib.SMA(close, timeperiod=period)
                df_features[f'EMA_{period}'] = talib.EMA(close, timeperiod=period)
            macd, macdsignal, macdhist = talib.MACD(close, fastperiod=12, slowperiod=26, signalperiod=9)
            df_features['MACD'] = macd; df_features['MACD_signal'] = macdsignal; df_features['MACD_hist'] = macdhist
            df_features['SAR'] = talib.SAR(high, low, acceleration=0.02, maximum=0.2)
            df_features['ADX_14'] = talib.ADX(high, low, close, timeperiod=14)
            for period in [9, 14, 21]: df_features[f'RSI_{period}'] = talib.RSI(close, timeperiod=period)
            slowk, slowd = talib.STOCH(high, low, close, fastk_period=14, slowk_period=3, slowk_matype=0, slowd_period=3, slowd_matype=0)
            df_features['STOCH_K'] = slowk; df_features['STOCH_D'] = slowd
            df_features['WILLR_14'] = talib.WILLR(high, low, close, timeperiod=14)
            for period in [5, 10, 20]: df_features[f'ROC_{period}'] = talib.ROC(close, timeperiod=period)
            df_features['CCI_14'] = talib.CCI(high, low, close, timeperiod=14)
            upper20, middle20, lower20 = talib.BBANDS(close, timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)
            df_features['BB_upper_20'] = upper20; df_features['BB_middle_20'] = middle20; df_features['BB_lower_20'] = lower20
            middle20_safe = np.where(middle20 == 0, np.nan, middle20)
            bb_width = (upper20 - lower20) / middle20_safe
            df_features['BB_width_20'] = np.nan_to_num(bb_width, nan=0.0)
            for period in [7, 14, 21]: df_features[f'ATR_{period}'] = talib.ATR(high, low, close, timeperiod=period)
            df_features['OBV'] = talib.OBV(close, volume)
            df_features['CMF_20'] = talib.ADOSC(high, low, close, volume, fastperiod=3, slowperiod=10)
            vol_sma_20 = talib.SMA(volume, timeperiod=20)
            vol_sma_20_safe = np.where(vol_sma_20 == 0, np.nan, vol_sma_20)
            df_features['volume_SMA_20'] = vol_sma_20
            df_features['volume_ratio'] = np.nan_to_num(volume / vol_sma_20_safe, nan=1.0)
            df_features['DOJI'] = talib.CDLDOJI(open_price, high, low, close) / 100
            df_features['ENGULFING'] = talib.CDLENGULFING(open_price, high, low, close) / 100
            df_features['HAMMER'] = talib.CDLHAMMER(open_price, high, low, close) / 100
            df_features['SHOOTING_STAR'] = talib.CDLSHOOTINGSTAR(open_price, high, low, close) / 100
        except Exception as e: feature_logger.error(f"Error calculating TA-Lib indicators: {e}", exc_info=True)

    if not TALIB_AVAILABLE:
        feature_logger.warning("Calculating subset of features manually (TA-Lib unavailable)")
        for period in [5, 10, 20]:
            df_features[f'SMA_{period}'] = df_features['close'].rolling(window=period).mean()
            df_features[f'EMA_{period}'] = df_features['close'].ewm(span=period, adjust=False).mean()
            # df_features[f'ROC_{period}'] = df_features['close'].pct_change(periods=period) * 100 # Replaced by return_xm
        delta = df_features['close'].diff(); gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean().replace(0, 1e-10); rs = gain / loss
        df_features['RSI_14'] = 100 - (100 / (1 + rs))

    # Derived Features (always calculated)
    for period in [14, 30]:
        roll_max = df_features['close'].rolling(window=period).max()
        roll_min = df_features['close'].rolling(window=period).min()
        # Ensure the range isn't too small to avoid extreme ratios
        min_range = df_features['close'].mean() * 0.001  # 0.1% of mean price as minimum range
        denominator = np.maximum(roll_max - roll_min, min_range)
        # Relative position is a ratio - apply Fisher transform to make more normal
        rel_pos = np.clip((df_features['close'] - roll_min) / denominator, 0.01, 0.99)
        # Apply logit transform (inverse of sigmoid) to normalize the distribution
        df_features[f'rel_position_{period}'] = np.log(rel_pos / (1 - rel_pos))

    for period in [20, 50]:
        if f'SMA_{period}' in df_features.columns:
            sma = df_features[f'SMA_{period}']
            close = df_features['close']
            # Apply log transform to price-to-SMA ratio to make distribution more normal
            ratio = np.clip(close / np.maximum(sma, close.mean() * 0.01), 0.5, 2.0)
            df_features[f'price_dist_SMA_{period}'] = np.log(ratio)

    # Intraday return - with log transform for better distribution
    open_prices = df_features['open'].replace(0, np.nan)
    close_prices = df_features['close']
    # Use a reasonable minimum price to avoid extreme ratios
    min_open = close_prices.mean() * 0.01  # 1% of mean close
    open_safe = np.maximum(open_prices.fillna(close_prices), min_open)
    # Apply log to price ratio for better distribution
    ratio = np.clip(close_prices / open_safe, 0.5, 2.0)
    df_features['intraday_return'] = np.log(ratio)
    
    # Log return with improved handling
    prev_close = df_features['close'].shift(1).replace(0, np.nan)
    current_close = df_features['close']
    # Avoid division by very small values
    ratio = np.clip(current_close / np.maximum(prev_close.fillna(current_close), current_close.mean() * 0.01), 0.5, 2.0)
    df_features['log_return'] = np.log(ratio)
    
    # Volatility with robust handling - apply log transform to make more normal
    log_return_clipped = np.clip(df_features['log_return'], -0.2, 0.2)
    vol = log_return_clipped.rolling(window=14).std().fillna(0)
    # Log transform volatility for more normal distribution
    df_features['volatility_14d'] = np.log1p(vol * 100)  # log1p avoids issues with zero values

    # Add Crypto Features
    df_features = add_crypto_features(df_features)

    # Final NaN fill (important after all calculations)
    # Fill with 0 as done in V6 for most features
    cols_to_fill = df_features.columns.difference(required_cols)
    for col in cols_to_fill:
        if df_features[col].isnull().any():
            df_features[col].fillna(0, inplace=True)

    # Drop original OHLCV + volume columns? No, keep them for potential analysis/state
    # df_features = df_features.drop(columns=required_cols)

    return df_features
# --- End V6 Feature Logic ---


class TradingSystem:
    """
    V7 Trading System: Integrates GRU Predictor and SAC Agent.
    Generates features, uses GRU for price/uncertainty prediction, 
    and SAC for trading decisions.
    """
    
    # V7.2 Add system logger
    _logger = logging.getLogger(__name__)
    if not _logger.hasHandlers():
        _logger.setLevel(logging.INFO)
        handler = logging.StreamHandler(sys.stdout)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        handler.setFormatter(formatter)
        _logger.addHandler(handler)
        _logger.propagate = False
    
    def __init__(self, gru_model: CryptoGRUModel = None, sac_agent: SimplifiedSACTradingAgent = None, gru_lookback=60):
        """
        Initialize the TradingSystem.

        Args:
            gru_model (CryptoGRUModel, optional): Pre-initialized GRU model. Defaults to None.
            sac_agent (SACTradingAgent, optional): Pre-initialized SAC agent. Defaults to None.
                                                    If None, a default agent with state_dim=2 is created.
            gru_lookback (int): Number of time steps for GRU input sequence.
        """
        self._logger.info("Initializing V7 Trading System...")
        self.gru_model = gru_model or CryptoGRUModel() # Use provided or create new
        self.sac_agent = sac_agent or SimplifiedSACTradingAgent(state_dim=2) # V7 state: [pred_return, uncertainty]
        self.gru_lookback = gru_lookback
        self.feature_scaler = None
        self.y_scaler = None
        self.last_prediction = None
        self.experiences = [] # Store (state, action, reward, next_state, done)
        # V7.23: Add storage for scaled feature indices
        self._momentum_feature_idx = None
        self._volatility_feature_idx = None
        # Initialize scalers from GRU model if available and loaded
        if self.gru_model and self.gru_model.is_loaded:
            self.feature_scaler = self.gru_model.feature_scaler
            self.y_scaler = self.gru_model.y_scaler
            _logger.info("Scalers initialized from pre-loaded GRU model.")
            # V7.23: Attempt to set indices if scaler loaded
            self._set_feature_indices()

    # V7.23: Helper method to find and store feature indices
    def _set_feature_indices(self):
        """Finds and stores the column indices for momentum and volatility features."""
        if self.feature_scaler is None:
            self._logger.warning("Feature scaler not available. Cannot set feature indices.")
            self._momentum_feature_idx = None
            self._volatility_feature_idx = None
            return False

        if not hasattr(self.feature_scaler, 'feature_names_in_'):
            # Fallback for older sklearn versions or different scaler types
            if hasattr(self.feature_scaler, 'n_features_in_'):
                 self._logger.warning("Feature scaler lacks 'feature_names_in_'. Cannot reliably find named features. State scaling may fail.")
                 # Cannot set indices reliably without names
                 self._momentum_feature_idx = None
                 self._volatility_feature_idx = None
                 return False
            else:
                 self._logger.error("Feature scaler lacks both 'feature_names_in_' and 'n_features_in_'. Cannot determine feature indices.")
                 self._momentum_feature_idx = None
                 self._volatility_feature_idx = None
                 return False

        feature_columns = list(self.feature_scaler.feature_names_in_)
        try:
            # Use the correct feature names: 'return_5m' and 'volatility_14d'
            self._momentum_feature_idx = feature_columns.index('return_5m')
            self._volatility_feature_idx = feature_columns.index('volatility_14d')
            self._logger.info(f"Successfully set feature indices: momentum_5m={self._momentum_feature_idx}, volatility_14d={self._volatility_feature_idx}")
            return True
        except ValueError as e:
            self._logger.error(f"Could not find 'return_5m' or 'volatility_14d' in scaler's feature columns: {feature_columns}. Error: {e}. State construction will likely fail.")
            self._momentum_feature_idx = None
            self._volatility_feature_idx = None
            return False

    # V7-V6 Update: Extract features and predict return/uncertainty for SAC state
    def _extract_features_and_predict(self, data_df_full, current_idx):
        """
        Extracts feature sequence and uses the trained GRU model to predict 
        the next period's return and associated confidence/uncertainty.

        Args:
            data_df_full: Full DataFrame with OHLCV and potentially pre-calculated features.
            current_idx: The current index in the DataFrame.

        Returns:
            Tuple(np.array, float, float) or (None, None, None):
                - features_sequence_scaled: Scaled NumPy array (sequence_length, num_features)
                - predicted_return: GRU predicted return for the next period.
                - uncertainty_sigma: Uncertainty estimate for the next period.
        """
        # --- Step 1: Extract Raw Feature Sequence ---
        # Determine the slice needed for feature calculation (V6 style)
        # Needs enough history for longest V6 feature calc + sequence length
        # Assume V6 features require ~100 periods + sequence length
        required_feature_history = 100 + self.gru_lookback # Conservative estimate
        if current_idx < required_feature_history - 1:
            # feature_logger.debug(f"Not enough history at idx {current_idx} for feature calc + lookback ({required_feature_history})")
            return None, None, None

        start_idx_calc = max(0, current_idx - required_feature_history + 1)
        end_idx_calc = current_idx + 1 # Include current index
        data_slice_for_calc = data_df_full.iloc[start_idx_calc:end_idx_calc]

        # --- Step 2: Calculate All V6 Features ---
        try:
            df_with_features = calculate_v6_features(data_slice_for_calc)
            if df_with_features.empty or len(df_with_features) < self.gru_lookback:
                feature_logger.warning(f"Feature calculation failed or insufficient length at idx {current_idx}")
                return None, None, None
        except Exception as e:
            feature_logger.error(f"Error calculating V6 features at idx {current_idx}: {e}", exc_info=False)
            return None, None, None

        # --- Step 3: Select Final Sequence & Scale ---
        # The sequence ends at the *current* index (current_idx)
        feature_sequence_unscaled_df = df_with_features.iloc[-self.gru_lookback:]

        # Check if feature scaler exists (should after GRU training/loading)
        if self.gru_model is None or self.feature_scaler is None:
            feature_logger.error(f"GRU model or feature scaler not available at idx {current_idx}. Cannot scale features.")
            return None, None, None

        try:
            # Ensure columns match scaler's expectations
            expected_features = self.feature_scaler.feature_names_in_
            # Select and reorder columns
            feature_sequence_aligned_df = feature_sequence_unscaled_df[expected_features]
            
            # Scale the sequence - Pass the DataFrame directly
            # V7 Warning Fix: Pass DataFrame to transform, not .values
            features_sequence_scaled_array = self.feature_scaler.transform(feature_sequence_aligned_df)
            
            # Ensure output is numpy array for reshaping
            if isinstance(features_sequence_scaled_array, pd.DataFrame):
                 features_sequence_scaled = features_sequence_scaled_array.values
            else:
                 features_sequence_scaled = features_sequence_scaled_array # Assume it's already numpy
                 
            if features_sequence_scaled.shape[0] != self.gru_lookback:
                 feature_logger.error(f"Scaled sequence shape incorrect at idx {current_idx}: {features_sequence_scaled.shape}")
                 return None, None, None
        except KeyError as e:
             feature_logger.error(f"Missing expected feature columns for scaler at idx {current_idx}: {e}")
             feature_logger.error(f" Available columns: {feature_sequence_unscaled_df.columns.tolist()}")
             feature_logger.error(f" Expected columns: {expected_features}")
             return None, None, None
        except Exception as e:
            feature_logger.error(f"Error scaling feature sequence at idx {current_idx}: {e}", exc_info=False)
            return None, None, None

        # --- Step 4: Predict Return and Uncertainty using GRU ---
        # Perform efficient single-step prediction and MC dropout here
        try:
            # Reshape for model prediction (1, seq_len, n_features)
            model_input = features_sequence_scaled.reshape(1, self.gru_lookback, features_sequence_scaled.shape[1])

            # 4a. Get standard prediction (scaled price)
            pred_scaled = self.gru_model.model.predict(model_input, verbose=0)[0, 0]

            # 4b. Perform MC Dropout for uncertainty
            n_mc_samples = 30 # Use a reasonable number of samples
            mc_preds_scaled_list = []
            
            # Define the prediction step with training=True
            @tf.function
            def mc_predict_step_single(batch):
                return self.gru_model.model(batch, training=True) # Enable dropout

            for _ in range(n_mc_samples):
                 mc_pred_scaled = mc_predict_step_single(tf.constant(model_input, dtype=tf.float32)).numpy()
                 mc_preds_scaled_list.append(mc_pred_scaled[0, 0]) # Get the scalar prediction
            
            mc_preds_scaled_array = np.array(mc_preds_scaled_list)
            mc_std_scaled = np.std(mc_preds_scaled_array)

            # 4c. Unscale prediction and uncertainty
            if self.y_scaler is None:
                 feature_logger.error(f"Y-scaler not available at idx {current_idx}, cannot unscale prediction/uncertainty.")
                 return None, None, None
                 
            pred_unscaled = self.y_scaler.inverse_transform([[pred_scaled]])[0, 0]
            mc_unscaled_std_dev = 0.0 # Default
            try:
                 # Use scaler's data range
                 if hasattr(self.y_scaler, 'data_min_') and hasattr(self.y_scaler, 'data_max_'):
                     data_range = self.y_scaler.data_max_[0] - self.y_scaler.data_min_[0]
                     if data_range > 1e-9:
                         mc_unscaled_std_dev = mc_std_scaled * data_range
                     else: feature_logger.warning("Scaler data range is near zero for unscaling std dev.")
                 else: feature_logger.warning("y_scaler missing data_min_/data_max_ for unscaling std dev.")
            except Exception as e:
                 feature_logger.error(f"Error unscaling MC std dev at idx {current_idx}: {e}")

            # 4d. Calculate predicted return
            # Get the *actual* close price at the *current* index to calculate return
            # Use the unscaled features dataframe for this
            start_price = feature_sequence_unscaled_df.iloc[-1]['close']
            epsilon = 1e-9
            if abs(start_price) > epsilon:
                 predicted_return = (pred_unscaled / start_price) - 1.0
            else:
                 predicted_return = 0.0
                 feature_logger.warning(f"Start price is near zero at idx {current_idx}, cannot calculate return.")

            # Return the SCALED features, predicted return, and UNCERTAINTY SIGMA
            return features_sequence_scaled, predicted_return, mc_unscaled_std_dev

        except Exception as e:
            feature_logger.error(f"Error during GRU prediction/evaluation step at idx {current_idx}: {e}", exc_info=False)
            return None, None, None

    # V7-V6 Update: Rename and modify to reflect price target
    def _preprocess_data_for_gru_training(self, df_data: pd.DataFrame, prediction_horizon: int) -> Tuple[pd.DataFrame, pd.Series, pd.Series]:
        """
        Preprocess data specifically for GRU model training (V6 Style - Price Target).
        Calculates V6 features, defines future price target, and aligns data.
        Scaling and sequence creation will happen *within* train_gru.
        
        Args:
            df_data: Raw OHLCV DataFrame
            prediction_horizon: How many steps ahead to predict the price.
        
        Returns:
            Tuple of (features DataFrame, future_price_target Series, start_price Series)
            Returns (None, None, None) if processing fails.
        """
        if df_data is None or df_data.empty:
            logging.error("Empty/None dataframe provided to _preprocess_data_for_gru_training")
            return None, None, None
        
        logging.info(f"Preprocessing data for GRU training (Target: Future Price, Horizon: {prediction_horizon})...")
        
        # --- Calculate V6 Features --- 
        logging.debug(f"Calculating V6 features for data shape {df_data.shape}")
        try:
            features_df = calculate_v6_features(df_data.copy()) # Use a copy
            if features_df is None or features_df.empty:
                logging.error("Feature calculation returned empty DataFrame.")
                return None, None, None
        except Exception as e:
            logging.error(f"Error calculating V6 features: {e}", exc_info=True)
            return None, None, None
        
        # --- Define Target (Future Price) & Start Price --- 
        try:
            # Target is the closing price `prediction_horizon` steps into the future
            future_price_target_ser = features_df['close'].shift(-prediction_horizon)
            
            # Start price is the current closing price (aligned with features)
            start_price_ser = features_df['close'] 
            
            # --- Align Features, Target, and Start Price --- 
            # Drop rows where the future target is NaN (typically at the end)
            common_index = future_price_target_ser.dropna().index
            
            # Ensure indices exist in all structures before slicing
            common_index = common_index.intersection(features_df.index).intersection(start_price_ser.index)
            
            if common_index.empty:
                 logging.error("No common index found after aligning features and future price target.")
                 return None, None, None
                 
            features_aligned_df = features_df.loc[common_index]
            target_aligned_ser = future_price_target_ser.loc[common_index]
            start_price_aligned_ser = start_price_ser.loc[common_index]
            
        except KeyError as e:
             logging.error(f"Missing 'close' column for target/start price calculation: {e}")
             return None, None, None
        except Exception as e:
            logging.error(f"Error defining/aligning target and start price: {e}", exc_info=True)
            return None, None, None
        
        # Final check for NaN values introduced during alignment or feature calc
        if features_aligned_df.isna().any().any():
            nan_cols = features_aligned_df.columns[features_aligned_df.isna().any()].tolist()
            logging.warning(f"NaN values detected in final aligned feature columns: {nan_cols}. Dropping rows.")
            features_aligned_df = features_aligned_df.dropna()
            # Realign target and start price
            common_index = features_aligned_df.index
            target_aligned_ser = target_aligned_ser.loc[common_index]
            start_price_aligned_ser = start_price_aligned_ser.loc[common_index]

        # --- Final Shape Check ---
        if not (len(features_aligned_df) == len(target_aligned_ser) == len(start_price_aligned_ser)):
            logging.error(f"Final shape mismatch after NaN handling: Features={len(features_aligned_df)}, Target={len(target_aligned_ser)}, StartPrice={len(start_price_aligned_ser)}")
            return None, None, None
        if features_aligned_df.empty:
             logging.error("Preprocessing resulted in empty features DataFrame after alignment.")
             return None, None, None
            
        logging.info(f"Preprocessing complete. Features: {features_aligned_df.shape}, Target: {target_aligned_ser.shape}, Start Price: {start_price_aligned_ser.shape}")
        return features_aligned_df, target_aligned_ser, start_price_aligned_ser

    def generate_trading_experiences(self, val_data, transaction_cost=0.00001, prediction_horizon=1):
        """
        V7-V6 Update: Use predicted return and uncertainty sigma for SAC state.
        V7 Efficiency Update: Pre-compute GRU predictions/uncertainty for the whole validation set.
        """
        print(f"Generating trading experiences for SAC agent on {len(val_data)} validation data points...")
        
        # --- Pre-computation Step ---
        if self.gru_model is None or not (self.gru_model.is_trained or self.gru_model.is_loaded):
            logging.error("Cannot generate experiences: GRU model not ready.")
            return []
        if self.feature_scaler is None or self.y_scaler is None:
            logging.error("Cannot generate experiences: Scalers not loaded/trained.")
            return []

        logging.info("Preprocessing validation data for experience generation...")
        val_features_df, val_target_price_ser, val_start_price_ser = self._preprocess_data_for_gru_training(val_data, prediction_horizon)
        if val_features_df is None:
            logging.error("Failed to preprocess validation data for experiences.")
            return []
            
        logging.info("Scaling validation features and targets...")
        try:
            val_features_scaled = self.feature_scaler.transform(val_features_df.select_dtypes(include=np.number).fillna(0))
            val_target_scaled = self.y_scaler.transform(val_target_price_ser.fillna(0).values.reshape(-1, 1))
        except Exception as e:
            logging.error(f"Error scaling validation data for experiences: {e}", exc_info=True)
            return []
            
        logging.info("Creating validation sequences...")
        try:
            val_start_price_aligned_ser = val_start_price_ser.loc[val_features_df.index] # Align before sequencing
            X_val, y_val_scaled_seq, y_start_price_val_seq = create_sequences_v2(
                pd.DataFrame(val_features_scaled, index=val_features_df.index), 
                pd.Series(val_target_scaled.flatten(), index=val_features_df.index), 
                val_start_price_aligned_ser,
                self.gru_lookback
            )
            if X_val is None or y_val_scaled_seq is None or y_start_price_val_seq is None:
                logging.error("Sequence creation failed for validation data.")
                return []
        except Exception as e:
            logging.error(f"Error creating validation sequences: {e}", exc_info=True)
            return []

        logging.info(f"Pre-computing GRU predictions/uncertainty for {len(X_val)} validation sequences...")
        try:
            eval_results = self.gru_model.evaluate(X_val, y_val_scaled_seq, y_start_price_val_seq, n_mc_samples=30)
            if eval_results is None:
                logging.error("GRU evaluate failed on validation sequences.")
                return []
            all_pred_returns = eval_results['pred_percent_change']
            all_uncertainties = eval_results['mc_unscaled_std_dev']
        except Exception as e:
            logging.error(f"Error during pre-computation evaluate call: {e}", exc_info=True)
            return []

        # --- V7.15 START: Extract Momentum and Volatility for Experience Generation --- 
        # Similar to backtest logic, extract features aligned with sequences
        num_sequences = len(all_pred_returns) # Use num_sequences from evaluate results
        all_momentum_5 = np.zeros(num_sequences)
        all_volatility_20 = np.zeros(num_sequences)
        try:
            # Check required columns in the *unscaled* preprocessed features df
            required_state_cols = ['return_5m', 'volatility_14d'] 
            if not all(col in val_features_df.columns for col in required_state_cols):
                missing_cols = [col for col in required_state_cols if col not in val_features_df.columns]
                logging.error(f"Missing required state columns in val_features_df: {missing_cols}")
                return []

            # Align features with the sequences
            if len(val_features_df) >= self.gru_lookback - 1 + num_sequences:
                aligned_feature_indices = val_features_df.index[self.gru_lookback - 1 : self.gru_lookback - 1 + num_sequences]
                aligned_features_for_state = val_features_df.loc[aligned_feature_indices]
                
                all_momentum_5 = aligned_features_for_state['return_5m'].values
                all_volatility_20 = aligned_features_for_state['volatility_14d'].values 
                
                if len(all_momentum_5) != num_sequences or len(all_volatility_20) != num_sequences:
                     logging.error(f"Exp Gen: Length mismatch extracting state features: Mom5({len(all_momentum_5)}), Vol20({len(all_volatility_20)}) vs NumSeq({num_sequences})")
                     return []
            else:
                 logging.error("Exp Gen: Length mismatch: val_features_df too short to extract aligned state features.")
                 return []
        except KeyError as e:
             logging.error(f"Exp Gen: KeyError extracting state features: {e}")
             return []
        except Exception as e:
             logging.error(f"Exp Gen: Error extracting state features: {e}", exc_info=True)
             return []
        # --- V7.15 END: Extract Momentum and Volatility --- 
        
        logging.info("GRU pre-computation finished. Generating experiences loop...")
        # --- End Pre-computation ---
        
        experiences = []
        current_position = 0.0 # Position *before* taking action at step i
        
        # Loop through the *results* (length = num_sequences)
        # num_sequences = len(all_pred_returns) # Already defined
        if num_sequences <= 1:
            logging.warning("Not enough sequences generated from validation data to create experiences.")
            return []
            
        # Need original close prices aligned with the sequences for reward calculation
        # The i-th sequence corresponds to the original data ending at index i + gru_lookback - 1
        # The prediction is for the step *after* this sequence
        # The relevant close prices for reward are at original indices [i + gru_lookback] and [i + gru_lookback + 1]
        original_close_prices = val_data['close'] # Use original validation data
        
        for i in tqdm(range(num_sequences - 1), desc="Generating Experiences"): # Iterate up to second-to-last sequence result
            # --- V7.15 START: Construct 5D state s_t ---
            # V7.23 Get SCALED momentum/volatility for state t
            pred_return_t = all_pred_returns[i]
            uncertainty_t = all_uncertainties[i]
            momentum_5_t_scaled = all_momentum_5[i] # Use scaled value
            volatility_20_t_scaled = all_volatility_20[i] # Use scaled value
            
            # Calculate z_proxy using position *before* action (current_position)
            # V7.23: Use SCALED volatility for consistency within state.
            z_proxy_t = current_position * volatility_20_t_scaled 
            state = np.array([pred_return_t, uncertainty_t, z_proxy_t, momentum_5_t_scaled, volatility_20_t_scaled], dtype=np.float32)
            
            # Handle potential NaNs/Infs
            if np.any(np.isnan(state)) or np.any(np.isinf(state)):
                 logging.warning(f"NaN/Inf in state at step {i}. Replacing with 0. State: {state}")
                 state = np.nan_to_num(state, nan=0.0, posinf=0.0, neginf=0.0)
            # --- V7.15 END: Construct 5D state s_t ---

            # Get action a_t (unpack action and log_prob)
            action_tuple = self.sac_agent.get_action(state, deterministic=False)
            if isinstance(action_tuple, (tuple, list)) and len(action_tuple) > 0:
                 action = action_tuple[0] # Action is the first element
                 # Ensure action is scalar if action_dim is 1
                 if self.sac_agent.action_dim == 1 and isinstance(action, (np.ndarray, list)):
                     action = action[0] 
            else:
                 logging.error(f"SAC agent get_action did not return expected tuple at step {i}. Got: {action_tuple}. Skipping experience.")
                 continue # Skip this experience if action format is wrong
            
            # --- V7.15 START: Construct 5D next_state s_{t+1} ---
            # V7.23 Get SCALED momentum/volatility for state t+1
            pred_return_t1 = all_pred_returns[i+1]
            uncertainty_t1 = all_uncertainties[i+1]
            momentum_5_t1_scaled = all_momentum_5[i+1] # Use scaled value
            volatility_20_t1_scaled = all_volatility_20[i+1] # Use scaled value
            
            # Calculate z_proxy using the *action* taken (action is position for next step)
            # V7.23: Use SCALED volatility for consistency
            z_proxy_t1 = action * volatility_20_t1_scaled # Use action, not current_position
            next_state = np.array([pred_return_t1, uncertainty_t1, z_proxy_t1, momentum_5_t1_scaled, volatility_20_t1_scaled], dtype=np.float32)
            
            # Handle potential NaNs/Infs
            if np.any(np.isnan(next_state)) or np.any(np.isinf(next_state)):
                 logging.warning(f"NaN/Inf in next_state at step {i}. Replacing with 0. State: {next_state}")
                 next_state = np.nan_to_num(next_state, nan=0.0, posinf=0.0, neginf=0.0)
            # --- V7.15 END: Construct 5D next_state s_{t+1} ---

            # Calculate actual return for reward r_t
            # Map sequence index 'i' back to the original dataframe index
            original_idx_t = i + self.gru_lookback # Index of the *last* element of the sequence i
            original_idx_t_plus_1 = original_idx_t + 1 # Index for the next closing price
            
            try:
                # Use original_close_prices Series which should have the same index as val_data
                close_t = original_close_prices.iloc[original_idx_t]
                close_t1 = original_close_prices.iloc[original_idx_t_plus_1]
                
                if close_t != 0:
                    actual_return = (close_t1 / close_t) - 1.0
                else:
                    actual_return = 0.0 # Assign 0 return if start price is zero
            except IndexError:
                logging.warning(f"IndexError accessing original close prices for reward at step {i} (original indices {original_idx_t}, {original_idx_t_plus_1}). Skipping experience.")
                continue
            except Exception as e:
                logging.error(f"Error accessing original close prices for reward at step {i}: {e}")
                continue
                 
            # Reward calculation (unchanged, uses action)
            reward = action * actual_return - transaction_cost * abs(action - current_position)

            # Done flag (only True for the very last possible transition)
            done = (i == num_sequences - 2) 
            # Store the experience (state, action, reward, next_state, done)
            experiences.append((state, [action], reward, next_state, float(done))) # Ensure action is in a list for consistency if needed by buffer
            
            # Update current_position for the *next* iteration's z_proxy_t calculation
            current_position = action 
            
        print(f"Generated {len(experiences)} experiences efficiently.")
        return experiences

    def train_sac(self, val_data, epochs=100, batch_size=256, transaction_cost=0.00001, 
                  generate_new_experiences_on_epoch=False, # Default is now False
                  prediction_horizon=1): # Need prediction horizon for experience gen
        """
        Train SAC agent using experiences generated from historical data.
        (Code largely unchanged, relies on updated generate_trading_experiences)
        """
        print(f"Starting SAC training for {epochs} epochs...")
        if not generate_new_experiences_on_epoch:
            print("Generating initial experiences..."); experiences = self.generate_trading_experiences(val_data, transaction_cost)
            print(f"Adding {len(experiences)} experiences to replay buffer...")
            for state, action, reward, next_state, done in tqdm(experiences, desc="Filling Buffer"):
                 action_to_add = action[0] if isinstance(action, (list, np.ndarray)) and len(action) == 1 else action
                 # V7.16 Fix: Use agent's store_transition method
                 # self.sac_agent.buffer.add(state, action_to_add, reward, next_state, done)
                 self.sac_agent.store_transition(state, action_to_add, reward, next_state, done)
            print("Initial buffer fill complete.")
        metrics_history = []
        for epoch in tqdm(range(epochs), desc="SAC Training Epochs"):
            if generate_new_experiences_on_epoch:
                experiences = self.generate_trading_experiences(val_data, transaction_cost, prediction_horizon) # Pass prediction_horizon
                for state, action, reward, next_state, done in experiences:
                      action_to_add = action[0] if isinstance(action, (list, np.ndarray)) and len(action) == 1 else action
                      self.sac_agent.buffer.add(state, action_to_add, reward, next_state, done)
            # V7.2 Revert: Use SAC agent default min buffer size
            # V7.17 Fix: Check buffer size using agent's buffer_counter
            # if len(self.sac_agent.buffer) >= self.sac_agent.min_buffer_size if hasattr(self.sac_agent, 'min_buffer_size') else batch_size:
            if self.sac_agent.buffer_counter >= (self.sac_agent.min_buffer_size if hasattr(self.sac_agent, 'min_buffer_size') else batch_size):
                # V7.17 Fix: train returns tuple (actor_loss, critic_loss) or None
                loss_tuple = self.sac_agent.train(batch_size) 
                if loss_tuple:
                    actor_loss, critic_loss = loss_tuple
                    metrics = {'actor_loss': actor_loss, 'critic_loss': critic_loss}
                    metrics_history.append(metrics)
                    # Print metrics every 10 epochs instead of 100
                    if metrics and epoch % 10 == 0: print(f"Epoch {epoch}/{epochs}: {metrics}")
                else:
                    # Handle case where train() returns None (e.g., not enough data yet)
                    pass # Or log a warning if needed
            elif epoch == 0: 
                 min_buff = self.sac_agent.min_buffer_size if hasattr(self.sac_agent, 'min_buffer_size') else batch_size
                 # V7.16 Fix: Use buffer_counter in log message
                 print(f"Epoch {epoch}: Buffer size ({self.sac_agent.buffer_counter}) < min ({min_buff}). Skipping.")
        print("SAC training complete."); return metrics_history

    def backtest_simple(self, data, transaction_cost=0.00001):
        """
        V7-V6 Update: Use predicted return and confidence for SAC state.
        """
        print(f"Starting simple backtest on {len(data)} data points...")
        portfolio_value = 1.0; current_position = 0.0
        portfolio_values = [portfolio_value]; positions_history = [current_position]
        predictions = []; uncertainties = [] # Store pred_return and confidence
        
        for i in tqdm(range(len(data) - 1), desc="Simple Backtest"):
            # V7-V6 Update: Use _extract_features_and_predict
            _, pred_return, uncertainty_sigma = self._extract_features_and_predict(data, i)
            # V7-V6 Sigma Update: Check uncertainty_sigma
            if pred_return is None or uncertainty_sigma is None:
                continue 
            
            predictions.append(pred_return)
            # V7-V6 Sigma Update: Store uncertainty_sigma
            uncertainties.append(uncertainty_sigma)

            # State for SAC: [predicted_return, mc_unscaled_std_dev]
            # V7-V6 Sigma Update: Use uncertainty_sigma
            state = np.array([pred_return, uncertainty_sigma], dtype=np.float32)
            
            action = self.sac_agent.get_action(state, deterministic=True)[0]

            tx_cost_fraction = transaction_cost * abs(action - current_position)
            
            # Use actual close prices for backtest return calc
            if data.iloc[i]['close'] != 0:
                price_return = (data.iloc[i+1]['close'] / data.iloc[i]['close']) - 1.0 
            else:
                price_return = 0
                
            position_return = current_position * price_return
            portfolio_value *= (1.0 + position_return) * (1.0 - tx_cost_fraction)
            positions_history.append(action); portfolio_values.append(portfolio_value)
            current_position = action
        
        # Calculate metrics
        portfolio_values = np.array(portfolio_values)
        returns = np.diff(portfolio_values) / portfolio_values[:-1] # Calculate portfolio returns
        returns = np.nan_to_num(returns)
        
        sharpe_ratio = 0.0
        if np.std(returns) > 1e-9:
            # Assume daily data for annualization? Needs correct frequency.
            # Let's assume minutes, T = minutes in a year
            T = 365 * 24 * 60
            sharpe_ratio = (np.mean(returns) / np.std(returns)) * np.sqrt(T) 
            
        max_drawdown = self._calculate_max_drawdown(portfolio_values)
        results = {
            'final_value': portfolio_values[-1],
            'total_return': portfolio_values[-1] - 1.0, 
            'sharpe_ratio': sharpe_ratio, 
            'max_drawdown': max_drawdown,
            'positions': positions_history, 
            'portfolio_values': portfolio_values, 
            'predictions': predictions, # Stored predicted returns
            'uncertainties': uncertainties # Stored confidence scores
        }
        print("Simple backtest complete.")
        print(f" Final Value: {results['final_value']:.4f}, Sharpe: {results['sharpe_ratio']:.4f}, Max Drawdown: {results['max_drawdown']:.4f}")
        return results

    def _calculate_max_drawdown(self, portfolio_values):
        """Calculate maximum drawdown from portfolio values"""
        if len(portfolio_values) < 2:
            return 0.0
        values = np.array(portfolio_values)
        running_max = np.maximum.accumulate(values)
        running_max[running_max == 0] = 1.0 # Avoid division by zero
        drawdown = (running_max - values) / running_max
        return np.max(drawdown)

    def save(self, path):
        """Save the integrated trading system (GRU model/scalers and SAC agent)."""
        print(f"Saving trading system to {path}...")
        os.makedirs(path, exist_ok=True)
        gru_path = os.path.join(path, "gru_model")
        sac_path = os.path.join(path, "sac_agent")
        os.makedirs(gru_path, exist_ok=True)
        os.makedirs(sac_path, exist_ok=True)

        if self.gru_model:
            # V7-V6 Update: Use CryptoGRUModel save
            self.gru_model.save(gru_path) # CryptoGRUModel.save handles model+scalers
        if self.sac_agent:
            self.sac_agent.save(sac_path)
        print("Trading system saved.")

    def load(self, path):
        """Load the integrated trading system (GRU model/scalers and SAC agent)."""
        print(f"Loading trading system from {path}...")
        gru_path = os.path.join(path, "gru_model")
        sac_path = os.path.join(path, "sac_agent")
        models_loaded = True
        try:
            if os.path.isdir(gru_path):
                # V7-V6 Update: Instantiate and load CryptoGRUModel
                self.gru_model = CryptoGRUModel()
                if self.gru_model.load(gru_path):
                    self.feature_scaler = self.gru_model.feature_scaler
                    self.y_scaler = self.gru_model.y_scaler
                    # V7.23: Set feature indices after loading scaler
                    if self._set_feature_indices():
                         gru_ok = True
                    else:
                         print("Warning: Failed to set feature indices after loading GRU scaler.")
                         gru_ok = False # Mark as failed if indices can't be set
                else:
                     print("Warning: Failed to load GRU model/scalers.")
            else: print(f"Warning: GRU model directory not found: {gru_path}")
            if os.path.isdir(sac_path):
                 # V7.23 Use correct state dim when loading
                 self.sac_agent = SimplifiedSACTradingAgent(state_dim=5)
                 if self.sac_agent.load(sac_path):
                     sac_ok = True
                 else:
                     print("Warning: Failed to load SAC agent.")
            else: print(f"Warning: SAC agent directory not found: {sac_path}")
            # V7.23 Check if both models loaded successfully
            models_loaded = gru_ok and sac_ok
            status = 'successful' if models_loaded else ('partially successful' if gru_ok or sac_ok else 'failed')
            print(f"Trading system loading {status}.")
        except Exception as e:
            print(f"An error occurred during loading: {e}")
            # Ensure flags reflect failure
            models_loaded = False

        return models_loaded # Return overall success status

    # V7-V6 Update: Adapt GRU training call for price regression
    def train_gru(self, train_data: pd.DataFrame, val_data: pd.DataFrame, 
                  prediction_horizon: int, epochs=20, batch_size=32, 
                  patience=10, # Revert to V6 defaults
                  model_save_dir='models/gru_predictor_trained'):
        """V7 Adaptation: Train GRU model (V6 style - price prediction). """
        print(f"--- Starting GRU Training Pipeline (V6 Adaptation) ---")
        print(f"Preprocessing training data (Target: Future Price, Horizon: {prediction_horizon})...")
        # V7-V6 Update: Use renamed preprocessing function
        train_features_df, train_target_price_ser, train_start_price_ser = self._preprocess_data_for_gru_training(train_data, prediction_horizon)
        if train_features_df is None: 
            print("ERROR: Failed to preprocess training data for GRU.")
            return None
        
        print(f"Preprocessing validation data (Target: Future Price, Horizon: {prediction_horizon})...")
        # V7-V6 Update: Use renamed preprocessing function
        val_features_df, val_target_price_ser, val_start_price_ser = self._preprocess_data_for_gru_training(val_data, prediction_horizon)
        if val_features_df is None: 
            print("ERROR: Failed to preprocess validation data for GRU.")
            return None
        
        # --- V7-V6 Update: Add Scaling Step --- 
        logging.info("Scaling features and target price...")
        try:
            # Feature Scaling (Use StandardScaler as in V6 preproc example, but make flexible)
            # Consider making scaler type a parameter if needed
            feature_scaler = StandardScaler()
            # Ensure features are numeric and handle potential NaNs before scaling
            train_features_numeric = train_features_df.select_dtypes(include=np.number).fillna(0)
            val_features_numeric = val_features_df.select_dtypes(include=np.number).fillna(0)
            
            # Log columns used for scaling
            logging.info(f"Feature columns for scaling: {train_features_numeric.columns.tolist()}")
            
            train_features_scaled = feature_scaler.fit_transform(train_features_numeric)
            val_features_scaled = feature_scaler.transform(val_features_numeric)
            self.feature_scaler = feature_scaler # Store fitted scaler
            logging.info(f"Feature scaling complete. Scaler type: {type(feature_scaler)}")
            
            # Target Scaling (Use MinMaxScaler as in V6)
            y_scaler = MinMaxScaler() # Default range (0, 1) is fine for prices
            # Handle potential NaNs in target before scaling
            train_target_numeric = train_target_price_ser.fillna(0).values.reshape(-1, 1)
            val_target_numeric = val_target_price_ser.fillna(0).values.reshape(-1, 1)
            
            train_target_scaled = y_scaler.fit_transform(train_target_numeric)
            val_target_scaled = y_scaler.transform(val_target_numeric)
            self.y_scaler = y_scaler # Store fitted scaler
            logging.info(f"Target price scaling complete. Scaler type: {type(y_scaler)}")
            
            # V7.23: Set feature indices now that scaler is fitted
            if not self._set_feature_indices():
                logging.error("Failed to set feature indices after scaling. Cannot proceed with GRU training.")
                return None # Exit if indices aren't found
            
        except Exception as e:
             logging.error(f"Error during scaling: {e}", exc_info=True)
             return None # Corrected indentation
        # --- End Scaling Step ---

        # --- V7-V6 Update: Add Sequence Creation Step ---
        logging.info("Creating sequences...")
        try: # Added missing except block below
             # Pass scaled features as DataFrame to preserve column info if needed by create_sequences_v2?
             # Or just pass the numpy array? Let's assume numpy is fine for V6 logic.
             # Also need to pass the UNscaled start price series correctly aligned.

             # Align start price series with the scaled data length before sequencing
             train_start_price_aligned_ser = train_start_price_ser.loc[train_features_df.index] # Index from before scaling
             val_start_price_aligned_ser = val_start_price_ser.loc[val_features_df.index]

             X_train, y_train_scaled, y_start_price_train = create_sequences_v2(
                 pd.DataFrame(train_features_scaled, index=train_features_df.index), # Pass features with index
                 pd.Series(train_target_scaled.flatten(), index=train_features_df.index), # Pass target with index
                 train_start_price_aligned_ser, # Pass aligned start price Series
                 self.gru_lookback
             )
             X_val, y_val_scaled, y_start_price_val = create_sequences_v2(
                 pd.DataFrame(val_features_scaled, index=val_features_df.index),
                 pd.Series(val_target_scaled.flatten(), index=val_features_df.index),
                 val_start_price_aligned_ser,
                 self.gru_lookback
             )

             if X_train is None or X_val is None or y_train_scaled is None or y_val_scaled is None:
                 logging.error("Sequence creation failed. Returned None.")
                 return None

        except Exception as e: # Added missing except block
             logging.error(f"Error creating sequences: {e}", exc_info=True) # Corrected indentation
             return None # Corrected indentation
        # --- End Sequence Creation ---

        # Debug info about the final data shapes for training
        print(f"Data ready for GRU training:")
        print(f"  X_train shape: {X_train.shape}")
        print(f"  y_train_scaled shape: {y_train_scaled.shape}")
        print(f"  X_val shape: {X_val.shape}")
        print(f"  y_val_scaled shape: {y_val_scaled.shape}")
        
        # Initialize GRU model if needed
        if self.gru_model is None:
            # V7-V6 Update: Use CryptoGRUModel
            self.gru_model = CryptoGRUModel()
        
        # --- V7-V6 Update: Call CryptoGRUModel.train --- 
        print(f"Training GRU model (V6 Adaptation) for {epochs} epochs...")
        history = self.gru_model.train(
            X_train, y_train_scaled, # Pass sequences and scaled targets
            X_val, y_val_scaled,
            feature_scaler=self.feature_scaler, # Pass fitted scalers
            y_scaler=self.y_scaler,
            epochs=epochs, 
            batch_size=batch_size,
            patience=patience,
            model_save_dir=model_save_dir # Pass save dir
            # Removed LR args, handled within model train
        )
        # --- End Updated Call --- 
        
        if history is not None:
            print(f"GRU model training complete. Model and scalers saved to {model_save_dir}.")
            # Store scalers in TradingSystem instance after successful training
            self.feature_scaler = self.gru_model.feature_scaler
            self.y_scaler = self.gru_model.y_scaler
            
            # Optional: Plot training history
            try:
                 history_plot_path = os.path.join(model_save_dir, 'gru_training_history.png')
                 self.gru_model.plot_training_history(history, save_path=history_plot_path)
            except Exception as plot_e:
                 logging.warning(f"Could not plot GRU training history: {plot_e}")
                 
            return history
        else:
            print("GRU model training failed.")
            return None


class ExtendedBacktester:
    """
    Enhanced backtesting framework for GRU+SAC integration.
    V7 Efficiency Update: Pre-computes GRU outputs for faster backtesting.
    """
    def __init__(self, trading_system: TradingSystem, initial_capital=10000.0, transaction_cost=0.0001, instrument_label="Unknown Instrument"):
        print("Initializing ExtendedBacktester...")
        self.trading_system = trading_system
        self.initial_capital = initial_capital
        self.transaction_cost = transaction_cost
        self.instrument_label = instrument_label # Store instrument label
        self.portfolio_values = []
        self.positions = []
        self.predictions = [] # Will store predicted returns used for state
        self.uncertainties_used_in_state = [] # Will store uncertainty sigmas used for state
        self.all_precomputed_uncertainties = None # Stores the full array from evaluate()
        self.actions = [] # Actions taken by SAC
        self.actual_returns = [] # Actual price returns observed
        self.trade_counts = 0
        self.trade_history = []
        self.timestamps = [] # Timestamps for portfolio/action logging
        self.gru_lookback = self.trading_system.gru_lookback # Get from TradingSystem
        self.buy_hold_values = None # Initialize B&H value storage
        # Add storage for GRU prediction plot data
        self.prediction_timestamps = None
        self.predicted_prices = None
        self.true_prices_for_pred = None
        # self.uncertainties already exists for uncertainty sigma
        print("ExtendedBacktester initialized.")

    def backtest(self, test_data, verbose=True, prediction_horizon=1):
        """V7 Efficiency Update: Pre-compute GRU outputs."""
        if not isinstance(test_data, pd.DataFrame) or 'close' not in test_data.columns:
            raise ValueError("test_data must be a pandas DataFrame with a 'close' column.")
        if self.trading_system.gru_model is None or not (self.trading_system.gru_model.is_trained or self.trading_system.gru_model.is_loaded):
             raise ValueError("Cannot backtest: GRU model not ready.")
        if self.trading_system.feature_scaler is None or self.trading_system.y_scaler is None:
             raise ValueError("Cannot backtest: Scalers not loaded/trained.")
        if not self.trading_system.sac_agent:
            raise ValueError("Cannot backtest: SAC Agent not initialized.")

        print("--- Pre-computing GRU outputs for Test Set ---")

        # 1. Preprocess test data (Features, Target Price, Start Price)
        feature_logger.info("Preprocessing test data...") # Use feature_logger
        test_features_df, test_target_price_ser, test_start_price_ser = self.trading_system._preprocess_data_for_gru_training(test_data, prediction_horizon)
        if test_features_df is None:
             feature_logger.error("Failed to preprocess test data.") # Use feature_logger
             return None

        # 2. Scale Features and Target
        feature_logger.info("Scaling test features and targets...") # Use feature_logger
        try:
             test_features_scaled = self.trading_system.feature_scaler.transform(test_features_df.select_dtypes(include=np.number).fillna(0))
             test_target_scaled = self.trading_system.y_scaler.transform(test_target_price_ser.fillna(0).values.reshape(-1, 1))
        except Exception as e:
             feature_logger.error(f"Error scaling test data: {e}", exc_info=True) # Use feature_logger
             return None

        # 3. Create Sequences
        feature_logger.info("Creating test sequences...") # Use feature_logger
        try:
            test_start_price_aligned_ser = test_start_price_ser.loc[test_features_df.index] # Align before sequencing
            X_test, y_test_scaled_seq, y_start_price_test_seq = create_sequences_v2(
                 pd.DataFrame(test_features_scaled, index=test_features_df.index),
                 pd.Series(test_target_scaled.flatten(), index=test_features_df.index),
                 test_start_price_aligned_ser,
                 self.gru_lookback
             )
            if X_test is None or y_test_scaled_seq is None or y_start_price_test_seq is None:
                 feature_logger.error("Sequence creation failed for test data.") # Use feature_logger
                 return None
        except Exception as e:
             feature_logger.error(f"Error creating test sequences: {e}", exc_info=True) # Use feature_logger
             return None

        # 4. Run GRU Evaluate Once
        feature_logger.info(f"Pre-computing GRU predictions/uncertainty for {len(X_test)} test sequences...") # Use feature_logger
        try:
            eval_results = self.trading_system.gru_model.evaluate(X_test, y_test_scaled_seq, y_start_price_test_seq, n_mc_samples=30)
            if eval_results is None:
                 feature_logger.error("GRU evaluate failed on test sequences.") # Use feature_logger
                 return None
            all_pred_returns = eval_results['pred_percent_change']
            all_uncertainties = eval_results['mc_unscaled_std_dev']
            # Store data needed for GRU prediction plot
            self.predicted_prices = eval_results.get('predicted_unscaled_prices')
            self.true_prices_for_pred = eval_results.get('true_unscaled_prices')
            # V7 Fix: Store full uncertainty array separately
            self.all_precomputed_uncertainties = all_uncertainties # Store full array for plotting
            # self.uncertainties will store the ones used in the loop state below
            self.uncertainties_used_in_state = [] # Initialize loop list

            # Store the corresponding timestamps
            # The prediction results align with the *end* of the sequences
            # The index of test_start_price_aligned_ser aligns with the start of the window
            # So, the timestamp for the k-th prediction is at index k + gru_lookback - 1
            num_sequences_eval = len(eval_results['predicted_unscaled_prices']) # Get length from results
            if len(test_start_price_aligned_ser) >= self.gru_lookback - 1 + num_sequences_eval:
                 prediction_indices = test_start_price_aligned_ser.index[self.gru_lookback - 1 : self.gru_lookback - 1 + num_sequences_eval]
                 self.prediction_timestamps = prediction_indices.to_list()
            else:
                 feature_logger.error("Cannot extract prediction timestamps due to index length mismatch.")
                 self.prediction_timestamps = [] # Assign empty list on error
            
            # Validate stored prediction data
            if self.predicted_prices is None or self.true_prices_for_pred is None or self.uncertainties_used_in_state is None:
                feature_logger.error("Missing price/uncertainty data from GRU evaluate results.")
                return None
            if len(self.prediction_timestamps) != len(self.predicted_prices):
                 feature_logger.error("Prediction timestamp length mismatch.")
                 return None
                 
        except Exception as e:
             feature_logger.error(f"Error during pre-computation evaluate call on test data: {e}", exc_info=True) # Use feature_logger
             return None

        # --- Calculate num_sequences FIRST ---
        # Moved from below to be available for feature extraction
        num_sequences = len(all_pred_returns)
        if num_sequences <= 1:
            feature_logger.warning("Not enough sequences generated from test data to perform backtest.") # Use feature_logger
            return None

        # --- V7.13 START: Extract Momentum and Volatility Features ---
        # These features should be aligned with the *sequences* used for GRU prediction
        # Use test_features_df which is aligned with the sequences' start times
        try:
            # Check if required columns exist in the preprocessed features
            required_state_cols = ['return_5m', 'volatility_14d'] # Example columns used in state
            if not all(col in test_features_df.columns for col in required_state_cols):
                missing_cols = [col for col in required_state_cols if col not in test_features_df.columns]
                feature_logger.error(f"Missing required state columns in test_features_df: {missing_cols}")
                return None

            # Select the features corresponding to the *end* of each sequence
            # The evaluate results (all_pred_returns) align with the *output* of the sequences
            # We need features from the *last timestep* of the input sequence (index k + gru_lookback - 1)
            # Slice test_features_df to match the number of sequences evaluated
            if len(test_features_df) >= self.gru_lookback - 1 + num_sequences:
                aligned_feature_indices = test_features_df.index[self.gru_lookback - 1 : self.gru_lookback - 1 + num_sequences]
                aligned_features_for_state = test_features_df.loc[aligned_feature_indices]
                
                # V7.14 Use return_5m for momentum_5, volatility_14d for volatility_20 (adapt names if needed)
                all_momentum_5 = aligned_features_for_state['return_5m'].values
                all_volatility_20 = aligned_features_for_state['volatility_14d'].values # Assuming 'volatility_14d' is the correct column
                
                if len(all_momentum_5) != num_sequences or len(all_volatility_20) != num_sequences:
                     feature_logger.error(f"Length mismatch extracting state features: Mom5({len(all_momentum_5)}), Vol20({len(all_volatility_20)}) vs NumSeq({num_sequences})")
                     return None
            else:
                 feature_logger.error("Length mismatch: test_features_df too short to extract aligned state features.")
                 return None
                 
        except KeyError as e:
             feature_logger.error(f"KeyError extracting state features (momentum/volatility): {e}")
             return None
        except Exception as e:
             feature_logger.error(f"Error extracting state features (momentum/volatility): {e}", exc_info=True)
             return None
        # --- V7.13 END: Extract Momentum and Volatility Features ---

        feature_logger.info("--- GRU Pre-computation Complete. Starting Backtest Simulation ---") # Use feature_logger

        # --- Refined Timestamp Handling START (from original test_data before preprocessing) ---
        time_col_name = None
        initial_timestamp_for_loop = None
        
        # Use a fresh copy for timestamp handling to avoid modifying the original input df
        test_data_copy = test_data.copy()
        
        if 'timestamp' in test_data_copy.columns:
            time_col_name = 'timestamp'
            test_data_copy[time_col_name] = pd.to_datetime(test_data_copy[time_col_name]) 
            # Get timestamp corresponding to the *end* of the first sequence
            if len(test_data_copy) > self.gru_lookback:
                initial_timestamp_for_loop = test_data_copy.iloc[self.gru_lookback][time_col_name]
            else:
                raise ValueError("Test data too short for lookback after timestamp check.")
            print(f"Using '{time_col_name}' column for time.")
            
        elif isinstance(test_data_copy.index, pd.DatetimeIndex):
            feature_logger.info("Using DatetimeIndex for time.") # Changed from print
            original_index_name = test_data_copy.index.name
            if len(test_data_copy) > self.gru_lookback:
                initial_timestamp_for_loop = test_data_copy.index[self.gru_lookback]
            else:
                raise ValueError("Test data index too short for lookback after timestamp check.")
            
            test_data_copy = test_data_copy.reset_index() 
            
            # Find the name of the column created from the index
            if original_index_name and original_index_name in test_data_copy.columns: 
                time_col_name = original_index_name
            elif 'index' in test_data_copy.columns and pd.api.types.is_datetime64_any_dtype(test_data_copy['index']): 
                 time_col_name = 'index'
            else: 
                found_col = None
                for col in test_data_copy.columns:
                     if pd.api.types.is_datetime64_any_dtype(test_data_copy[col]):
                         found_col = col; break
                if found_col: time_col_name = found_col
                else: raise ValueError("Could not identify the reset index column containing timestamps.")
            print(f"Identified reset index time column as: '{time_col_name}'")
        else:
            raise ValueError("Could not find 'timestamp' column or DatetimeIndex.")
        # --- Refined Timestamp Handling END ---

        # --- Backtest Loop Initialization ---
        self.portfolio_values = [self.initial_capital]
        self.positions = [0.0] # Position held *before* action at step i
        self.predictions = [] 
        self.uncertainties_used_in_state = []
        self.actions = [] 
        self.actual_returns = [] 
        self.trade_history = []
        self.trade_counts = 0
        # Initialize timestamp list with the timestamp for the END of the first sequence period
        self.timestamps = [initial_timestamp_for_loop] 

        portfolio_value = self.initial_capital
        current_position = 0.0 # Position held *before* taking action at step i

        # --- Main Backtest Loop (Iterates over pre-computed results) ---
        # Use original close prices from the time-handled copy
        original_close_prices = test_data_copy['close'] 

        for i in tqdm(range(num_sequences - 1), desc="Backtest Simulation", disable=not verbose):
            # 1. Get state from pre-computed results
            pred_return = all_pred_returns[i]
            uncertainty_sigma = all_uncertainties[i]
            # V7.23 Get SCALED momentum and volatility for the current step
            momentum_5_scaled = all_momentum_5[i]
            volatility_20_scaled = all_volatility_20[i]

            # Calculate z-proxy (Position as proxy for risk aversion)
            # V7.23 Use SCALED volatility for consistency in state
            z_proxy = current_position * volatility_20_scaled

            # Construct 5D state using SCALED momentum/volatility
            state = np.array([
                pred_return,
                uncertainty_sigma,
                z_proxy,
                momentum_5_scaled,
                volatility_20_scaled
            ], dtype=np.float32)

            # Handle NaNs/Infs in state
            if np.any(np.isnan(state)) or np.any(np.isinf(state)):
                 feature_logger.warning(f"NaN/Inf detected in state at step {i}. Replacing with 0. State: {state}")
                 state = np.nan_to_num(state, nan=0.0, posinf=0.0, neginf=0.0)

            # 2. Get deterministic action from SAC agent
            action_tuple = self.trading_system.sac_agent.get_action(state, deterministic=True)
            # V7.14 Ensure unpacking handles potential tuple vs single value return if agent API changes
            if isinstance(action_tuple, (tuple, list)) and len(action_tuple) > 0:
                 raw_action = action_tuple[0] # Take the first element (the action)
            else:
                 # Fallback if the API returns just the action (for robustness)
                 raw_action = action_tuple 
                 feature_logger.warning("SAC agent get_action did not return a tuple as expected. Assuming single value is action.")
            
            # V7.14 Ensure raw_action is a scalar if action_dim is 1
            if self.trading_system.sac_agent.action_dim == 1 and isinstance(raw_action, (np.ndarray, list)):
                 raw_action = raw_action[0]

            # 3. Calculate PnL and portfolio value
            position_change = raw_action - current_position

            # Map sequence index 'i' back to original dataframe index for price lookup
            # Use the index from the time-handled copy (test_data_copy)
            original_idx_t = i + self.gru_lookback # End of sequence i / Start of action period
            original_idx_t_plus_1 = original_idx_t + 1 # End of action period

            try:
                # Use prices from the time-handled copy
                current_price = original_close_prices.iloc[original_idx_t]
                next_price = original_close_prices.iloc[original_idx_t_plus_1]
                if current_price != 0:
                    price_return = (next_price / current_price) - 1.0
                else:
                    price_return = 0.0
            except IndexError:
                 feature_logger.warning(f"IndexError accessing original close prices for PnL at step {i} (original indices {original_idx_t}, {original_idx_t_plus_1}). Ending backtest early.")
                 break
            except Exception as e:
                 feature_logger.error(f"Error accessing original close prices for PnL at step {i}: {e}")
                 break

            tx_cost = abs(position_change) * portfolio_value * self.transaction_cost
            # PnL based on position held *during* the period (current_position)
            position_pnl = current_position * price_return * portfolio_value
            new_portfolio_value = portfolio_value + position_pnl - tx_cost

            # 4. Store results & update state for next iteration
            self.predictions.append(pred_return)
            # V7 Fix: Append uncertainty used in state to the correct list
            self.uncertainties_used_in_state.append(uncertainty_sigma)
            self.actions.append(raw_action)
            self.actual_returns.append(price_return)
            self.portfolio_values.append(new_portfolio_value)
            self.positions.append(raw_action) # Store position *after* action
            try:
                 # Use .loc with the identified time_col_name from the time-handled copy
                 self.timestamps.append(test_data_copy.loc[original_idx_t_plus_1, time_col_name])
            except KeyError:
                 feature_logger.warning(f"Could not find timestamp for index {original_idx_t_plus_1}. Appending last known timestamp.")
                 self.timestamps.append(self.timestamps[-1])

            if abs(position_change) > 0.01:
                trade_timestamp = test_data_copy.loc[original_idx_t, time_col_name] # Timestamp when action is decided
                trade = {'timestamp': trade_timestamp, 'price': current_price, 'old_position': current_position,
                         'new_position': raw_action, 'prediction': pred_return, 'uncertainty': uncertainty_sigma,
                         'portfolio_value_before': portfolio_value, 'portfolio_value_after': new_portfolio_value, 'transaction_cost': tx_cost}
                self.trade_history.append(trade)
                self.trade_counts += 1

            portfolio_value = new_portfolio_value
            current_position = raw_action # Position for the *next* period

            if portfolio_value <= 0:
                print(f"Portfolio zeroed at step {i}. Stopping.")
                # Adjust fill length based on remaining sequences
                fill_len = (num_sequences - 1) - (i + 1)
                self.portfolio_values.extend([0.0] * fill_len)
                self.positions.extend([current_position] * fill_len) # Fill with last position
                self.predictions.extend([0.0] * fill_len)
                # V7 Fix: Extend the correct list
                self.uncertainties_used_in_state.extend([0.0] * fill_len)
                self.actions.extend([current_position] * fill_len)
                self.actual_returns.extend([0.0] * fill_len)
                # Extend timestamps carefully using the time-handled copy
                last_valid_orig_idx = original_idx_t_plus_1
                if last_valid_orig_idx + fill_len <= len(test_data_copy) - 1:
                     remaining_timestamps = test_data_copy.loc[last_valid_orig_idx + 1 : last_valid_orig_idx + fill_len, time_col_name].tolist()
                     self.timestamps.extend(remaining_timestamps)
                else:
                     # If fill_len exceeds available data, add NaT or last timestamp
                     available_count = len(test_data_copy) - 1 - last_valid_orig_idx
                     if available_count > 0:
                        remaining_timestamps = test_data_copy.loc[last_valid_orig_idx + 1 : last_valid_orig_idx + available_count, time_col_name].tolist()
                        self.timestamps.extend(remaining_timestamps)
                     if fill_len - available_count > 0:
                        self.timestamps.extend([pd.NaT] * (fill_len - available_count))
                     break # Corrected indentation again
        # --- End Backtest Loop ---

        # --- Calculate Buy and Hold Benchmark --- 
        feature_logger.info("Calculating Buy and Hold benchmark...")
        try:
            # Use the original test_data before reset_index if applicable
            start_price_bh = test_data['close'].iloc[0]
            # Use the close price corresponding to the *last timestamp* used in the strategy backtest
            # The last timestamp added was for original_idx_t_plus_1 where i = num_sequences - 2
            last_strategy_idx = (num_sequences - 2) + self.gru_lookback + 1
            end_price_bh = test_data_copy['close'].iloc[last_strategy_idx]
            
            if pd.isna(start_price_bh) or start_price_bh <= 1e-9:
                feature_logger.warning("Could not calculate Buy & Hold: Invalid start price.")
                self.buy_hold_values = None
            else:
                initial_assets_bh = self.initial_capital / start_price_bh
                # Calculate B&H value for the same period as the strategy
                # Use close prices from the time-handled copy aligned with strategy timestamps
                relevant_close_prices = test_data_copy['close'].iloc[self.gru_lookback : last_strategy_idx + 1]
                self.buy_hold_values = initial_assets_bh * relevant_close_prices
                # Ensure the length matches the strategy portfolio values
                if len(self.buy_hold_values) != len(self.portfolio_values):
                    feature_logger.warning(f"Buy & Hold length mismatch ({len(self.buy_hold_values)}) vs Strategy ({len(self.portfolio_values)}). Adjusting B&H series.")
                    # Pad or truncate B&H to match strategy length (more robust approach needed if indices differ significantly)
                    bh_series = pd.Series(self.buy_hold_values, index=self.timestamps) # Use strategy timestamps
                    self.buy_hold_values = bh_series.reindex(self.timestamps).fillna(method='ffill').values
                else:
                     self.buy_hold_values = self.buy_hold_values.values # Convert to numpy
                     
                bh_final_value = self.buy_hold_values[-1]
                feature_logger.info(f"Buy & Hold Final Value: ${bh_final_value:.2f}")
        except Exception as e:
            feature_logger.error(f"Error calculating Buy & Hold benchmark: {e}", exc_info=True)
            self.buy_hold_values = None
        # --- End Buy and Hold Calculation --- 

        # Calculate & print metrics
        metrics = self._calculate_performance_metrics() # Now calculates B&H metrics internally
        metrics['trade_history'] = self.trade_history
        metrics['timestamps'] = self.timestamps # Already populated
        if verbose:
            print("\n--- Extended Backtest Complete ---")
            print(f"Final portfolio value: {metrics.get('final_value', 0):.2f}")
            # ... (print other metrics) ...
            print("---------------------------------")
        return metrics

    def _calculate_performance_metrics(self):
        if len(self.portfolio_values) < 2:
            return {'final_value': 0.0, 'initial_value': 0.0, 'total_return': 0.0, 'annual_return': 0.0, 'sharpe_ratio': 0.0, 'sortino_ratio': 0.0, 'volatility': 0.0, 'max_drawdown': 0.0, 'avg_position': 0.0, 'position_accuracy': 0.0, 'pred_accuracy': 0.0, 'prediction_rmse': 0.0, 'pred_return_corr': 0.0, 'pred_pos_corr': 0.0, 'unc_pos_corr': 0.0, 'total_trades': 0}
        portfolio_values = np.array(self.portfolio_values)
        positions = np.array(self.positions[1:])
        predictions = np.array(self.predictions)
        actual_returns = np.array(self.actual_returns)
        # V7 Fix: Use the uncertainties actually used in the state loop
        uncertainties_used = np.array(self.uncertainties_used_in_state)
        safe_portfolio_values = portfolio_values[:-1].copy()
        safe_portfolio_values[safe_portfolio_values == 0] = 1e-9
        period_returns = np.diff(portfolio_values) / safe_portfolio_values
        total_minutes = len(period_returns)
        minutes_per_year = 365 * 24 * 60
        if total_minutes == 0: 
            # Return default metrics if no periods
            return {'final_value': 0.0, 'initial_value': 0.0, 'total_return': 0.0, 'annual_return': 0.0, 'sharpe_ratio': 0.0, 'sortino_ratio': 0.0, 'volatility': 0.0, 'max_drawdown': 0.0, 'avg_position': 0.0, 'position_accuracy': 0.0, 'pred_accuracy': 0.0, 'prediction_rmse': 0.0, 'pred_return_corr': 0.0, 'pred_pos_corr': 0.0, 'unc_pos_corr': 0.0, 'total_trades': 0}
            
        final_value = portfolio_values[-1]
        initial_value = portfolio_values[0]
        total_return = (final_value / initial_value) - 1.0 if initial_value != 0 else 0.0
        annual_return = ((1.0 + total_return) ** (minutes_per_year / total_minutes)) - 1.0 if total_return > -1 else -1.0
        volatility = np.std(period_returns) * np.sqrt(minutes_per_year)
        mean_return = np.mean(period_returns)
        std_return = np.std(period_returns)
        sharpe_ratio = (mean_return / std_return) * np.sqrt(minutes_per_year) if std_return > 1e-9 else 0.0
        negative_returns = period_returns[period_returns < 0]
        downside_std = np.std(negative_returns) if len(negative_returns) > 0 else 0.0
        sortino_ratio = (mean_return / downside_std) * np.sqrt(minutes_per_year) if downside_std > 1e-9 else 0.0
        max_drawdown = self.trading_system._calculate_max_drawdown(portfolio_values)
        # V7 Fix: Use uncertainties_used here
        common_length = min(len(predictions), len(actual_returns), len(positions), len(uncertainties_used))
        if common_length > 0:
            aligned_predictions = predictions[:common_length]
            aligned_actual_returns = actual_returns[:common_length]
            aligned_positions = positions[:common_length]
            # V7 Fix: Use uncertainties_used here
            aligned_uncertainties = uncertainties_used[:common_length]
            avg_position = np.mean(np.abs(aligned_positions))
            non_zero_mask = aligned_positions != 0
            # Use nan_to_num on signs before comparison for robustness
            sign_pos = np.sign(np.nan_to_num(aligned_positions[non_zero_mask]))
            sign_ret = np.sign(np.nan_to_num(aligned_actual_returns[non_zero_mask]))
            position_accuracy = np.mean(sign_pos == sign_ret) if np.any(non_zero_mask) else 0.0
            pred_accuracy = np.mean(np.sign(aligned_predictions) == np.sign(aligned_actual_returns))
            prediction_rmse = np.sqrt(np.mean((aligned_predictions - aligned_actual_returns)**2))
            
            # Robust Correlation Calculation with warnings suppressed:
            with warnings.catch_warnings():
                # Ignore the specific RuntimeWarning from np.corrcoef with zero variance
                warnings.filterwarnings('ignore', r'invalid value encountered in divide', RuntimeWarning)
                
                pred_safe = np.nan_to_num(aligned_predictions)
                ret_safe = np.nan_to_num(aligned_actual_returns)
                pos_safe = np.nan_to_num(aligned_positions)
                unc_safe = np.nan_to_num(aligned_uncertainties)
                abs_pos_safe = np.nan_to_num(np.abs(aligned_positions))

                pred_return_corr = np.corrcoef(pred_safe, ret_safe)[0, 1] if np.std(pred_safe) > 1e-9 and np.std(ret_safe) > 1e-9 else 0.0
                pred_pos_corr = np.corrcoef(pred_safe, pos_safe)[0, 1] if np.std(pred_safe) > 1e-9 and np.std(pos_safe) > 1e-9 else 0.0
                unc_pos_corr = np.corrcoef(unc_safe, abs_pos_safe)[0, 1] if np.std(unc_safe) > 1e-9 and np.std(abs_pos_safe) > 1e-9 else 0.0

                # Explicitly handle potential NaNs from corrcoef if std dev check wasn't perfect
                if np.isnan(pred_return_corr): pred_return_corr = 0.0
                if np.isnan(pred_pos_corr): pred_pos_corr = 0.0
                if np.isnan(unc_pos_corr): unc_pos_corr = 0.0
        else: 
            avg_position, position_accuracy, pred_accuracy, prediction_rmse, pred_return_corr, pred_pos_corr, unc_pos_corr = (0.0,) * 7
        
        metrics = {
            'final_value': final_value, 'initial_value': initial_value, 'total_return': total_return,
            'annual_return': annual_return, 'sharpe_ratio': sharpe_ratio, 'sortino_ratio': sortino_ratio,
            'volatility': volatility, 'max_drawdown': max_drawdown, 'avg_position': avg_position,
            'position_accuracy': position_accuracy, 'pred_accuracy': pred_accuracy,
            'prediction_rmse': prediction_rmse, 'pred_return_corr': pred_return_corr,
            'pred_pos_corr': pred_pos_corr, 'unc_pos_corr': unc_pos_corr,
            'total_trades': self.trade_counts
        }

        # Add Buy & Hold metrics if available
        if self.buy_hold_values is not None and len(self.buy_hold_values) > 0:
            bh_final = self.buy_hold_values[-1]
            bh_initial = self.buy_hold_values[0] # Should correspond to initial capital approx.
            metrics['buy_hold_final_value'] = bh_final
            metrics['buy_hold_total_return'] = (bh_final / self.initial_capital) - 1.0 if self.initial_capital > 0 else 0.0
        else:
            metrics['buy_hold_final_value'] = None
            metrics['buy_hold_total_return'] = None

        for key, value in metrics.items():
             # Avoid rounding None
             if isinstance(value, (float, np.float_)):
                  metrics[key] = round(value, 6)
        return metrics

    # V7 Plot Update: Rewrite plot_results for 3 specific subplots
    def plot_results(self, save_path='backtest_results_v7.png', n_std_dev=1.5):
        """Generate and save a 3-panel plot: GRU Preds, SAC Actions, Portfolio Perf."""
        if plt is None:
            print("Matplotlib not installed. Skipping plotting.")
            return None

        # --- Log Timestamp Range Used for Plotting ---
        if self.timestamps:
            feature_logger.info(f"Plotting Results - Timestamp range: {min(self.timestamps)} to {max(self.timestamps)}")
        else:
            feature_logger.warning("Plotting Results - Timestamps list is empty.")

        # --- Construct Title Suffix --- 
        title_suffix = f" ({self.instrument_label})"
        if self.timestamps:
            try:
                start_ts_str = pd.to_datetime(min(self.timestamps)).strftime('%Y-%m-%d %H:%M')
                end_ts_str = pd.to_datetime(max(self.timestamps)).strftime('%Y-%m-%d %H:%M')
                title_suffix = f" ({self.instrument_label} | {start_ts_str} to {end_ts_str})"
            except Exception as e:
                 feature_logger.warning(f"Could not format timestamps for plot title: {e}")
        else:
            feature_logger.warning("Timestamps missing for plot title.")

        # --- Data Validation for Plotting ---
        valid_portfolio = self.portfolio_values is not None and len(self.portfolio_values) > 1
        valid_bh = self.buy_hold_values is not None and len(self.buy_hold_values) == len(self.portfolio_values)
        valid_actions = self.actions is not None and len(self.actions) == len(self.portfolio_values) -1 # Actions map to periods
        # V7 Fix: Check all_precomputed_uncertainties for plotting
        valid_gru_preds = (self.predicted_prices is not None and
                           self.true_prices_for_pred is not None and
                           self.all_precomputed_uncertainties is not None and # Check the correct variable
                           self.prediction_timestamps is not None and 
                           len(self.prediction_timestamps) == len(self.predicted_prices))
        
        # Timestamps for portfolio/action plots (start from the first action's effect)
        portfolio_time_axis = self.timestamps # Should align with portfolio_values
        
        # Timestamps for GRU predictions (align with the prediction time)
        gru_time_axis = pd.to_datetime(self.prediction_timestamps) if valid_gru_preds else None
        
        # Determine the number of plots needed
        n_plots = sum([valid_gru_preds, valid_actions, valid_portfolio])
        if n_plots == 0:
             print("Warning: No valid data available for any plots.")
             return None

        fig, axs = plt.subplots(n_plots, 1, figsize=(15, 6 * n_plots), sharex=True)
        # Ensure axs is always an array, even if n_plots is 1
        if n_plots == 1:
            axs = [axs]
        plot_idx = 0
        
        # --- Subplot 1: GRU Price Prediction vs Actual --- 
        if valid_gru_preds:
            ax = axs[plot_idx]
            # Align true prices with prediction timestamps
            true_prices = pd.Series(self.true_prices_for_pred, index=gru_time_axis)
            pred_prices = pd.Series(self.predicted_prices, index=gru_time_axis)
            # V7 Fix: Use all_precomputed_uncertainties for plotting
            if len(self.all_precomputed_uncertainties) != len(gru_time_axis):
                 print(f"Warning: Uncertainty length ({len(self.all_precomputed_uncertainties)}) doesn't match time axis ({len(gru_time_axis)}) for plotting.")
                 # Attempt to align or skip plotting uncertainty
                 uncertainty = None # Skip fill_between if lengths mismatch
            else:
                 uncertainty = pd.Series(self.all_precomputed_uncertainties, index=gru_time_axis)
            
            ax.plot(gru_time_axis, true_prices, label='Actual Price', color='#2ca02c', alpha=0.8, linewidth=1.5)
            ax.plot(gru_time_axis, pred_prices, label='Predicted Price (GRU)', color='#1f77b4', alpha=0.8, linewidth=1.5)
            
            # Only plot uncertainty band if uncertainty Series was created successfully
            if uncertainty is not None:
                lower_bound = pred_prices - n_std_dev * uncertainty
                upper_bound = pred_prices + n_std_dev * uncertainty
                ax.fill_between(gru_time_axis, lower_bound, upper_bound,
                                  color='#aec7e8', alpha=0.3,
                                  label=f'Prediction +/- {n_std_dev} Std Dev')
                               
            ax.set_title('GRU Price Prediction vs Actual' + title_suffix)
            ax.set_ylabel('Price')
            ax.legend()
            ax.grid(True)
            ax.ticklabel_format(style='plain', axis='y')
            plot_idx += 1
        else:
             print("Skipping GRU Prediction plot due to missing data.")

        # --- Subplot 2: SAC Agent Actions --- 
        if valid_actions:
            ax = axs[plot_idx]
            # Align actions with the correct timestamps (action taken at T determines position for T to T+1)
            action_time_axis = portfolio_time_axis[:-1] # Actions align with the start of the period
            if len(action_time_axis) != len(self.actions):
                 feature_logger.warning(f"Action timestamp length mismatch ({len(action_time_axis)}) vs actions ({len(self.actions)}). Skipping SAC plot.")
            else:
                 ax.plot(action_time_axis, self.actions, label='SAC Position Size', drawstyle='steps-post', color='#ff7f0e')
                 ax.set_title('SAC Agent Position Size (-1 to 1)' + title_suffix)
                 ax.set_ylabel('Position')
                 ax.set_ylim(-1.1, 1.1)
                 ax.grid(True)
                 ax.legend()
            plot_idx += 1
        else:
            print("Skipping SAC Action plot due to missing data.")

        # --- Subplot 3: Portfolio Performance vs Buy & Hold --- 
        if valid_portfolio:
            ax = axs[plot_idx]
            ax.plot(portfolio_time_axis, self.portfolio_values, label=f'Strategy Value: ${self.portfolio_values[-1]:,.2f}', color='blue')
            
            if valid_bh:
                 ax.plot(portfolio_time_axis, self.buy_hold_values, label=f'Buy & Hold Value: ${self.buy_hold_values[-1]:,.2f}', color='orange', linestyle='--')
            elif self.buy_hold_values is not None:
                 feature_logger.warning("Could not plot Buy & Hold line due to length mismatch.")
                 
            # Add trade markers
            if self.trade_history:
                buys = [(pd.to_datetime(t['timestamp']), t['price']) for t in self.trade_history if t['new_position'] > t['old_position']]
                sells = [(pd.to_datetime(t['timestamp']), t['price']) for t in self.trade_history if t['new_position'] < t['old_position']]
                if buys:
                    buy_times, buy_prices = zip(*buys)
                    # Find corresponding portfolio values at buy times for marker placement
                    portfolio_series = pd.Series(self.portfolio_values, index=pd.to_datetime(portfolio_time_axis))
                    buy_portfolio_values = portfolio_series.reindex(buy_times, method='ffill')
                    ax.plot(buy_times, buy_portfolio_values, '^', color='green', markersize=8, label='Buy Signal')
                if sells:
                    sell_times, sell_prices = zip(*sells)
                    portfolio_series = pd.Series(self.portfolio_values, index=pd.to_datetime(portfolio_time_axis))
                    sell_portfolio_values = portfolio_series.reindex(sell_times, method='ffill')
                    ax.plot(sell_times, sell_portfolio_values, 'v', color='red', markersize=8, label='Sell Signal')
            
            ax.set_title('Strategy Portfolio Value vs. Buy & Hold' + title_suffix)
            ax.set_ylabel('Portfolio Value ($)')
            ax.grid(True); ax.legend(); ax.ticklabel_format(style='plain', axis='y')
            plot_idx += 1
        else:
             print("Skipping Portfolio Performance plot due to missing data.")

        # Final adjustments
        axs[-1].set_xlabel('Time') # Add xlabel only to the bottom-most plot
        fig.autofmt_xdate()
        plt.tight_layout()
        
        if save_path:
            try: # Added missing except block below
                plt.savefig(save_path);
                print(f"Combined performance plot saved to {save_path}")
            except Exception as e: # Added missing except block
                 print(f"Error saving combined plot: {e}") # Corrected indentation
        # plt.show() # Optional: uncomment to display plot interactively
        plt.close();
        return fig

    def generate_performance_report(self, report_path="backtest_performance_report.md"):
        # ... (implementation remains the same, note update)
        metrics = self._calculate_performance_metrics()
        if not metrics: 
            feature_logger.error("Cannot generate report: No metrics calculated.")
            return None

        report = f"# GRU+SAC Backtesting Performance Report\n\n"
        report += f"Report generated on: {pd.Timestamp.now()}\n"
        # Ensure timestamps list is not empty before accessing
        if self.timestamps:
            # --- Log Timestamp Range Used for Report Header ---
            report_start_ts = self.timestamps[0]
            report_end_ts = self.timestamps[-1]
            feature_logger.info(f"Generating Report - Timestamp range: {report_start_ts} to {report_end_ts}")
            report += f"Data range: {self.timestamps[0]} to {self.timestamps[-1]}\n"
            report += f"Total duration: {self.timestamps[-1] - self.timestamps[0]}\n\n"
        else:
            feature_logger.warning("Generating Report - Timestamps list is empty.")
            report += "Data range: N/A\n"
            report += "Total duration: N/A\n\n"
             
        report += "## Strategy Performance Metrics\n\n"
        report += f"* **Initial capital:** ${metrics.get('initial_value', 0):,.2f}\n"
        report += f"* **Final portfolio value:** ${metrics.get('final_value', 0):,.2f}\n"
        report += f"* **Total return:** {metrics.get('total_return', 0)*100:.2f}%\n"
        report += f"* **Annualized return:** {metrics.get('annual_return', 0)*100:.2f}%\n"
        report += f"* **Sharpe ratio (annualized):** {metrics.get('sharpe_ratio', 0):.4f}\n"
        report += f"* **Sortino ratio (annualized):** {metrics.get('sortino_ratio', 0):.4f}\n"
        report += f"* **Volatility (annualized):** {metrics.get('volatility', 0)*100:.2f}%\n"
        report += f"* **Maximum drawdown:** {metrics.get('max_drawdown', 0)*100:.2f}%\n"
        report += f"* **Total trades:** {metrics.get('total_trades', 0)}\n"
        
        # Add Buy and Hold Performance section
        report += "\n## Buy and Hold Benchmark\n\n"
        bh_final_value = metrics.get('buy_hold_final_value')
        bh_total_return = metrics.get('buy_hold_total_return')
        if bh_final_value is not None and bh_total_return is not None:
            report += f"* **Final value (B&H):** ${bh_final_value:,.2f}\n"
            report += f"* **Total return (B&H):** {bh_total_return*100:.2f}%\n"
        else:
            report += "* *Buy and Hold benchmark could not be calculated.*\n"
            
        report += "\n## Position & Prediction Analysis\n\n"
        report += f"* **Average absolute position size:** {metrics.get('avg_position', 0):.4f}\n"
        report += f"* **Position sign accuracy vs return:** {metrics.get('position_accuracy', 0)*100:.2f}%\n"
        report += f"* **Prediction sign accuracy vs return:** {metrics.get('pred_accuracy', 0)*100:.2f}%\n"
        report += f"* **Prediction RMSE (on returns):** {metrics.get('prediction_rmse', 0):.6f}\n"
        report += "\n## Correlations\n\n"
        report += f"* **Prediction-Return correlation:** {metrics.get('pred_return_corr', 0):.4f}\n"
        report += f"* **Prediction-Position correlation:** {metrics.get('pred_pos_corr', 0):.4f}\n"
        report += f"* **Uncertainty-Position Size correlation:** {metrics.get('unc_pos_corr', 0):.4f}\n"
        report += "\n## Notes\n\n"
        report += f"* Transaction cost used: {self.transaction_cost * 100:.4f}% per position change value.\n"
        report += f"* GRU lookback period: {self.gru_lookback} minutes.\n"
        report += "* V6 features + return features used.\n"
        report += "* Uncertainty estimated via MC Dropout standard deviation.\n"
        # report += "* Uncertainty estimated directly by GRU second head.\n"
        if report_path:
            try:
                with open(report_path, "w") as f:
                    f.write(report)
                print(f"Performance report saved to {report_path}")
            except Exception as e:
                print(f"Error saving performance report: {e}")
        return report 