"""
Backtesting Engine.

Simulates trading strategy execution on historical test data, calculates 
performance metrics, and generates reports and plots.
"""

import os
import logging
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from sklearn.metrics import confusion_matrix, classification_report
import seaborn as sns
from typing import Dict, Any, Tuple, Optional, List
import inspect

# Import required components (use absolute paths)
from gru_sac_predictor.src.sac_agent import SACTradingAgent
from gru_sac_predictor.src.gru_model_handler import GRUModelHandler
from gru_sac_predictor.src.calibrator import Calibrator
from gru_sac_predictor.src.calibrator_vector import VectorCalibrator # Import needed for type hints
from gru_sac_predictor.src.metrics import edge_filtered_accuracy, calculate_sharpe_ratio, calculate_brier_score
# --- Import Metrics (Task 6.4) ---
try:
     from .metrics import edge_filtered_accuracy, calculate_sharpe_ratio, calculate_brier_score
except ImportError:
     logging.error("Failed to import metrics. Sharpe, Brier and Edge Acc will be missing.")
     # Define placeholders
     def edge_filtered_accuracy(*args, **kwargs): return np.nan, 0
     def calculate_sharpe_ratio(*args, **kwargs): return np.nan
     def calculate_brier_score(*args, **kwargs): return np.nan
# --- End Import --- #

logger = logging.getLogger(__name__)


def calculate_sharpe_ratio(returns, periods_per_year=252*24*60): # Default for 1-min data
    """Calculate annualized Sharpe ratio from a series of returns."""
    returns = pd.Series(returns)
    if returns.std() == 0 or np.isnan(returns.std()):
        return 0.0 if returns.mean() > 0 else (-np.inf if returns.mean() < 0 else 0.0)
    # Assuming risk-free rate is 0 for simplicity
    return np.sqrt(periods_per_year) * returns.mean() / returns.std()

def calculate_max_drawdown(equity_curve):
    """Calculate the maximum drawdown from an equity curve series."""
    equity_curve = pd.Series(equity_curve)
    rolling_max = equity_curve.cummax()
    drawdown = (equity_curve - rolling_max) / rolling_max
    max_drawdown = drawdown.min()
    return abs(max_drawdown) # Return positive value

class Backtester:
    """Runs the backtest simulation and generates results."""

    def __init__(self, config: dict, io_manager: Optional[Any] = None):
        """
        Initialize the Backtester.
        Args:
            config (dict): Pipeline configuration dictionary.
            io_manager (Optional[Any]): IOManager instance for saving results.
        """
        self.config = config
        self.io = io_manager
        self.env_cfg = config.get('environment', {})
        self.cal_cfg = config.get('calibration', {})
        self.sac_cfg = config.get('sac', {})
        
        self.initial_capital = self.env_cfg.get('initial_capital', 10000.0)
        self.transaction_cost = self.env_cfg.get('transaction_cost', 0.0005)
        
        # --- Store Rolling Calibration Config --- #
        self.rolling_cal_enabled = self.cal_cfg.get('rolling_enabled', False)
        self.recalibrate_every_n = self.cal_cfg.get('recalibrate_every_n', 5000)
        self.recalibration_window = self.cal_cfg.get('recalibration_window', 20000)
        self.coverage_alarm_enabled = self.cal_cfg.get('coverage_alarm_enabled', False)
        self.coverage_alarm_threshold_drop = self.cal_cfg.get('coverage_alarm_threshold_drop', 0.03)
        self.coverage_alarm_window = self.cal_cfg.get('coverage_alarm_window', 1000)
        # --- End --- #
        
        self.results_df: Optional[pd.DataFrame] = None
        self.metrics: Optional[Dict[str, Any]] = None
        
        logger.info("Backtester initialized.")
        logger.info(f"  Initial Capital: {self.initial_capital:.2f}")
        logger.info(f"  Transaction Cost: {self.transaction_cost*100:.4f}%")
        if self.rolling_cal_enabled:
            logger.info("  Rolling Calibration: Enabled")
            logger.info(f"    Recalibrate Every: {self.recalibrate_every_n} steps")
            logger.info(f"    Recalibration Window: {self.recalibration_window} steps")
            if self.coverage_alarm_enabled:
                 logger.info("    Coverage Alarm: Enabled")
                 logger.info(f"      Alarm Window: {self.coverage_alarm_window} steps")
                 logger.info(f"      Alarm Threshold Drop: {self.coverage_alarm_threshold_drop:.3f}")
        else:
             logger.info("  Rolling Calibration: Disabled")

        self.ece_recalibration_threshold = self.cal_cfg.get('ece_recalibration_threshold', 0.03)

    def run_backtest(
        self, 
        sac_agent_load_path: Optional[str],
        X_test_seq: np.ndarray,
        y_test_seq_dict: Dict[str, np.ndarray],
        test_indices: pd.Index, 
        gru_handler: GRUModelHandler,
        # --- Pass Calibrator Instances & Initial Params/Threshold --- #
        calibrator: Optional[Calibrator], # For Temp Scaling
        vector_calibrator: Optional[VectorCalibrator], # For Vector Scaling
        initial_optimal_T: Optional[float],
        initial_vector_params: Optional[np.ndarray],
        fold_edge_threshold: float, # Optimized or fixed threshold from validation
        # --- Pass Raw Predictions for Rolling Cal --- #
        p_raw_test: Optional[np.ndarray] = None,
        logits_test: Optional[np.ndarray] = None,
        # --- End Args --- #
        original_prices: Optional[pd.DataFrame] = None, # Pass DataFrame
        is_ternary: bool = False,
        fold_num: int = 0 # Added fold number
    ) -> Tuple[Optional[pd.DataFrame], Optional[Dict[str, Any]], Optional[pd.DataFrame]]:
        """
        Executes the backtesting simulation loop with optional rolling calibration.

        Args:
            sac_agent_load_path: Path to load the SAC agent.
            X_test_seq: Test set feature sequences.
            y_test_seq_dict: Dict of test set targets ('ret', 'dir'/'dir3').
            test_indices: Timestamps corresponding to the test sequences' targets.
            gru_handler: Instance to get GRU predictions (used only if raw preds not passed).
            calibrator: Instance for Temperature scaling (if used).
            vector_calibrator: Instance for Vector scaling (if used).
            initial_optimal_T: Optimal temperature from fold validation.
            initial_vector_params: Optimal vector params from fold validation.
            fold_edge_threshold: Edge threshold determined during validation (optimized or fixed).
            p_raw_test: Raw binary probabilities P(up) from GRU. Required if rolling_cal_enabled & !is_ternary.
            logits_test: Raw logits (N, 3) from GRU. Required if rolling_cal_enabled & is_ternary.
            original_prices: DataFrame with 'open', 'high', 'low', 'close', 'volume' aligned with test_indices.
            is_ternary: Flag indicating ternary classification.
            fold_num: Current fold number for logging.

        Returns:
            Tuple[Optional[pd.DataFrame], Optional[Dict[str, Any]], Optional[pd.DataFrame]]: 
                - DataFrame with detailed backtest results per step.
                - Dictionary containing calculated performance metrics.
                - DataFrame with periodic metrics logged during the backtest.
                Returns (None, None, None) if backtest cannot run.
        """
        logger.info(f"--- Fold {fold_num}: Starting Backtest Simulation --- ")

        # --- Input Validation --- #
        if X_test_seq is None or y_test_seq_dict is None or test_indices is None:
            logger.error(f"Fold {fold_num}: Test sequence data (X, y, indices) is missing. Cannot run backtest.")
            return None, None, None
        if gru_handler.model is None:
             logger.error(f"Fold {fold_num}: GRU model is not loaded in the handler. Cannot run backtest.")
             return None, None, None
        if original_prices is None:
            logger.warning(f"Fold {fold_num}: Original prices DataFrame not provided. Some metrics/plots may be unavailable.")
        if self.rolling_cal_enabled and not is_ternary and p_raw_test is None:
             logger.error(f"Fold {fold_num}: Rolling calibration enabled for binary case, but p_raw_test not provided.")
             return None, None, None
        if self.rolling_cal_enabled and is_ternary and logits_test is None:
             logger.error(f"Fold {fold_num}: Rolling calibration enabled for ternary case, but logits_test not provided.")
             return None, None, None
        # --- End Validation --- #
             
        # 1. Initialize SAC Agent & Load Weights
        # Ensure agent state dim matches the state construction below
        agent_state_dim = 5 # mu, sigma, edge, |mu|/sigma, position
        
        # --- Filter sac_cfg to only include valid SACTradingAgent.__init__ args --- 
        expected_args = set(inspect.signature(SACTradingAgent.__init__).parameters.keys())
        # Remove 'self' and potentially dims if handled separately
        expected_args -= {'self', 'state_dim', 'action_dim'}
        
        filtered_sac_cfg = {k: v for k, v in self.sac_cfg.items() if k in expected_args}
        logger.info(f"Filtered SAC config passed to agent: {filtered_sac_cfg.keys()}")
        # --- End Filter --- 
        
        agent = SACTradingAgent(
            state_dim=agent_state_dim,
            action_dim=1,
            # Pass FILTERED relevant SAC params from config
            **filtered_sac_cfg 
        )
        # Set the edge threshold for the agent (used in oracle seeding if enabled)
        agent.edge_threshold_config = fold_edge_threshold 
        
        if sac_agent_load_path and os.path.exists(sac_agent_load_path):
            logger.info(f"Fold {fold_num}: Loading SAC agent weights from: {sac_agent_load_path}")
            try:
                agent.load(sac_agent_load_path)
            except Exception as e:
                logger.error(f"Fold {fold_num}: Failed to load SAC agent weights from {sac_agent_load_path}: {e}. Proceeding with untrained agent.", exc_info=True)
        else:
            logger.warning(f"Fold {fold_num}: SAC agent load path not found or not specified ({sac_agent_load_path}). Proceeding with untrained agent.")
            
        # 2. Get GRU Predictions (if not passed for rolling cal) AND Pre-calculate Logits/Probs if needed
        mu_test_state, sigma_test_state = None, None 
        p_raw_fallback = None # For binary non-rolling case
        p_softmax_test = None # For ternary non-rolling/non-vector case (fallback)
        # Pre-calculate all test logits/raw_probs IF rolling cal is OFF to avoid step-by-step prediction
        all_logits_test_precalc = None
        all_p_raw_test_precalc = None

        if not self.rolling_cal_enabled:
            logger.info(f"Fold {fold_num}: Rolling cal disabled. Pre-calculating necessary GRU outputs for test set...")
            if is_ternary:
                 all_logits_test_precalc = gru_handler.predict_logits(X_test_seq)
                 # Also need mu for state construction
                 preds_test_mu = gru_handler.predict(X_test_seq) # Assumes predict returns mu as first item
                 if preds_test_mu is None: return None, None, None
                 mu_test_state = preds_test_mu[0].flatten()
                 sigma_test_state = np.ones_like(mu_test_state) * 0.01 # Placeholder sigma
                 if all_logits_test_precalc is None:
                      logger.error(f"Fold {fold_num}: Failed to pre-calculate logits needed for vector calibration.")
                      return None, None, None
                 else:
                      logger.info(f"Pre-calculated all_logits_test_precalc shape: {all_logits_test_precalc.shape}")
                      # Calculate softmax once if needed as fallback
                      import tensorflow as tf # Local import ok here
                      p_softmax_test = tf.nn.softmax(all_logits_test_precalc).numpy()
                      
            else: # Binary case, non-rolling
                 preds_test = gru_handler.predict(X_test_seq)
                 if preds_test is None or len(preds_test) < 3: return None, None, None # Error logged inside
                 mu_test_state = preds_test[0].flatten()
                 log_sigma_test_state = preds_test[1][:, 1].flatten() if preds_test[1].shape[-1] == 2 else np.log(preds_test[1].flatten() + 1e-9)
                 sigma_test_state = np.exp(log_sigma_test_state)
                 all_p_raw_test_precalc = preds_test[2].flatten() # Pre-calculate raw probs
                 p_raw_fallback = all_p_raw_test_precalc # Use pre-calculated as fallback
                 logger.info(f"Pre-calculated all_p_raw_test_precalc shape: {all_p_raw_test_precalc.shape}")

        else: # Rolling calibration IS enabled, use passed-in logits/p_raw
             logger.info(f"Fold {fold_num}: Rolling cal enabled. Using passed logits_test or p_raw_test.")
             # Still need mu/sigma for state construction
             if is_ternary:
                  # Logits were passed, get mu from separate prediction call
                  preds_test_mu = gru_handler.predict(X_test_seq) # Assumes predict returns mu as first item
                  if preds_test_mu is None: return None, None, None
                  mu_test_state = preds_test_mu[0].flatten()
                  sigma_test_state = np.ones_like(mu_test_state) * 0.01 # Placeholder sigma
             else: # Binary case, rolling
                  preds_test_state = gru_handler.predict(X_test_seq)
                  if preds_test_state is None or len(preds_test_state) < 2: # Need at least mu, sigma
                       logger.error(f"Fold {fold_num}: Failed to get GRU state predictions (mu, log_sigma) when p_raw was provided for rolling cal.")
                       return None, None, None
                  mu_test_state = preds_test_state[0].flatten()
                  log_sigma_test_state = preds_test_state[1][:, 1].flatten() if preds_test_state[1].shape[-1] == 2 else np.log(preds_test_state[1].flatten() + 1e-9)
                  sigma_test_state = np.exp(log_sigma_test_state)
             # p_raw_fallback is not needed when rolling cal is enabled, as p_raw_test is used directly later
             # p_softmax_test is not needed when rolling cal is enabled if vector cal is used

        # Extract actual returns and directions
        actual_ret_test = y_test_seq_dict.get('ret')
        dir_key = 'dir3' if is_ternary else 'dir'
        actual_dir_test = y_test_seq_dict.get(dir_key)
        if actual_ret_test is None or actual_dir_test is None:
             logger.error(f"Fold {fold_num}: Actual return ('ret') or direction ('{dir_key}') missing from y_test_seq_dict.")
             return None, None, None
             
        # Verify prediction lengths (ensure mu_test_state/sigma_test_state were calculated)
        n_test = len(X_test_seq)
        if mu_test_state is None or sigma_test_state is None:
             logger.error(f"Fold {fold_num}: Failed to extract mu/sigma state components. Cannot proceed.")
             return None, None, None
        if not (len(mu_test_state) == n_test and len(sigma_test_state) == n_test and 
                len(actual_ret_test) == n_test and len(actual_dir_test) == n_test and 
                len(test_indices) == n_test):
             logger.error(f"Fold {fold_num}: Length mismatch in test state predictions/targets/indices.")
             return None, None, None
             
        # --- Initialize Calibration State ---
        current_optimal_T = initial_optimal_T
        current_vector_params = initial_vector_params
        
        # --- Initialize Rolling Calibration State ---
        historical_data = [] # Stores tuples of (raw_pred, true_label)
        coverage_hits = []   # Buffer for coverage alarm hit rate calculation
        last_recalib_step = -1
        
        # --- Initialize Simulation State ---
        capital = self.initial_capital
        current_position = 0.0 # Starts neutral (-1 to 1)
        equity_curve = [capital]
        positions = [current_position]
        actions_taken = [0.0] # SAC agent's desired fractional position
        calibrated_probs_steps = [] # Store step-by-step calibrated prob
        edge_steps = []             # Store step-by-step edge
        pnl_steps = []
        trades_executed = [] # Store details of trades
        metrics_log = []     # Store periodic metrics
        step_correct_nonzero, step_count_nonzero = 0, 0
        step_abs_actions = []

        logger.info(f"Fold {fold_num}: Starting backtest simulation loop ({n_test} steps)...")
        for i in range(n_test):
            # --- Step i: Calibration ---
            p_cal_step = np.nan
            edge_step = np.nan
            
            if is_ternary and vector_calibrator is not None and current_vector_params is not None:
                # --- Use pre-calculated logits if available ---
                if all_logits_test_precalc is not None:
                    logit_step = all_logits_test_precalc[i:i+1]
                elif logits_test is not None: # Use passed logits if rolling cal enabled
                    logit_step = logits_test[i:i+1]
                else: # Should not happen if logic above is correct
                    logger.error(f"Step {i}: Logits not pre-calculated or passed. Cannot calibrate.")
                    logit_step = None
                # --- End Use pre-calculated ---

                if logit_step is None: 
                     logger.warning(f"Step {i}: Failed to get logits, using neutral prob."); p_cal_step=0.33; edge_step=0.0
                else:
                     # --- Corrected Call: Use calibrate() which uses internal W, b --- 
                     vector_calibrator.W = current_vector_params[:len(logit_step[0])]
                     vector_calibrator.b = current_vector_params[len(logit_step[0]):]
                     p_cal_step_all = vector_calibrator.calibrate(logit_step)
                     # --- End Correction ---
                     p_cal_step = p_cal_step_all[0, 2] # Prob(Up)
                     edge_step = p_cal_step_all[0, 2] - p_cal_step_all[0, 0] 

            elif not is_ternary and calibrator is not None and current_optimal_T is not None:
                # --- Use pre-calculated raw probs if available ---
                if all_p_raw_test_precalc is not None:
                    p_raw_step = all_p_raw_test_precalc[i]
                elif p_raw_test is not None: # Use passed raw probs if rolling cal enabled
                     p_raw_step = p_raw_test[i]
                else: # Should not happen
                     logger.error(f"Step {i}: Raw probs not pre-calculated or passed. Cannot calibrate.")
                     p_raw_step = None
                # --- End Use pre-calculated ---

                if p_raw_step is None: 
                     logger.warning(f"Step {i}: Failed to get raw prob, using neutral prob."); p_cal_step=0.5; edge_step=0.0
                else:
                     # --- Corrected Temp Scaling Call ---
                     # p_cal_step = calibrator.calibrate_with_T(p_raw_step, current_optimal_T)
                     calibrator.optimal_T = current_optimal_T # Ensure instance has the right T
                     p_cal_step = calibrator.calibrate(np.array([[p_raw_step]]))[0,0] # calibrate expects 2D
                     # --- End Correction ---
                     edge_step = 2 * p_cal_step - 1
            else:
                # Fallback if calibrator missing or not fitted - use raw prob? or neutral?
                if not is_ternary:
                     p_cal_step = p_raw_fallback[i] if p_raw_fallback is not None else 0.5 # Use pre-calculated raw if available
                     edge_step = 2 * p_cal_step - 1
                elif is_ternary and p_softmax_test is not None:
                     p_cal_step = p_softmax_test[i, 2] # Use raw softmax P(Up)
                     edge_step = p_softmax_test[i, 2] - p_softmax_test[i, 0] # Raw edge
                else: # Cannot determine calibrated prob or edge
                    p_cal_step = 0.5 if not is_ternary else 0.33
                    edge_step = 0.0 
                    
            calibrated_probs_steps.append(p_cal_step)
            edge_steps.append(edge_step)

            # --- Step i: State Construction ---
            z_score_step = np.abs(mu_test_state[i]) / (sigma_test_state[i] + 1e-9)
            state = np.array([
                mu_test_state[i], sigma_test_state[i], edge_step, z_score_step, current_position
            ], dtype=np.float32)

            # --- Step i: SAC Action ---
            sac_action = agent.get_action(state, deterministic=True)[0] 
            target_position = np.clip(sac_action, -1.0, 1.0) 
            
            # --- Step i: PnL Calculation ---
            step_actual_return = actual_ret_test[i]
            gross_pnl = current_position * capital * (np.exp(step_actual_return) - 1)
            trade = target_position - current_position
            cost = abs(trade) * capital * self.transaction_cost
            net_pnl = gross_pnl - cost
            capital += net_pnl
            
            # --- Step i: Store Results ---
            equity_curve.append(capital)
            positions.append(target_position) 
            actions_taken.append(sac_action)
            pnl_steps.append(net_pnl)
            if abs(trade) > 1e-6:
                 trades_executed.append({
                      'timestamp': test_indices[i], 'trade_size': trade, 'cost': cost,
                      'position_before': current_position, 'position_after': target_position
                 })

            # --- Step i: Update Position ---
            current_position = target_position 
            
            # --- Step i: Update Metrics Log Data ---
            if abs(current_position) > 1e-6: 
                step_count_nonzero += 1
                if (gross_pnl > 0 and current_position > 0) or \
                   (gross_pnl < 0 and current_position < 0) or \
                   (abs(gross_pnl) < 1e-9):
                    step_correct_nonzero += 1
            step_abs_actions.append(abs(sac_action))
            # Log metrics periodically (end of loop)

            # --- Step i: Rolling Calibration Update Logic ---
            recalibrate_now = False
            if self.rolling_cal_enabled:
                # Store data for potential refit
                true_label_step = actual_dir_test[i]
                # Use the raw prediction corresponding to the calibration method
                raw_pred_step = logits_test[i] if is_ternary else p_raw_test[i]
                # Only store if raw prediction was valid
                if raw_pred_step is not None: 
                    historical_data.append((raw_pred_step, true_label_step))
                
                # Select raw outputs and true labels for the window
                window_start_idx = max(0, i - self.recalibration_window + 1)
                window_labels_onehot = actual_dir_test[window_start_idx:i]
                window_logits = None
                window_raw_probs = None
                window_calibrated_probs = None # Store current calibrated probs for ECE
                if is_ternary and logits_test is not None:
                    window_logits = logits_test[window_start_idx:i]
                    if vector_calibrator is not None: # Ensure calibrator exists
                        window_calibrated_probs = vector_calibrator.calibrate(window_logits)
                elif not is_ternary and p_raw_test is not None:
                    window_raw_probs = p_raw_test[window_start_idx:i]
                    if calibrator is not None: # Ensure calibrator exists
                        window_calibrated_probs = calibrator.calibrate(window_raw_probs)
                
                # --- Check Coverage Alarm (Now using ECE - Revision 4) --- #
                trigger_recalibration = False
                if self.coverage_alarm_enabled and window_calibrated_probs is not None:
                    try:
                        # Prepare inputs for ECE calculation
                        if is_ternary:
                            probs_for_ece = window_calibrated_probs
                            labels_for_ece = window_labels_onehot
                        else: # Binary
                            # ECE helper expects N x K shape
                            p_binary_2class = np.vstack([1 - window_calibrated_probs, window_calibrated_probs]).T
                            # Convert labels to one-hot if not already
                            if window_labels_onehot.shape[1] == 1:
                                y_true_binary_indices = window_labels_onehot.flatten().astype(int)
                                labels_for_ece = tf.keras.utils.to_categorical(y_true_binary_indices, num_classes=2)
                            else:
                                labels_for_ece = window_labels_onehot
                            probs_for_ece = p_binary_2class

                            # Calculate ECE
                            if labels_for_ece.shape[0] > 0: # Ensure data exists
                                ece = self._calculate_ece(probs_for_ece, labels_for_ece)
                                logger.debug(f"Step {i}: Rolling ECE check: {ece:.4f} (Threshold: {self.ece_recalibration_threshold})")
                                if ece > self.ece_recalibration_threshold:
                                    trigger_recalibration = True
                                    logger.warning(f"Step {i}: ECE Coverage Alarm! ECE ({ece:.4f}) > Threshold ({self.ece_recalibration_threshold}). Triggering recalibration.")
                            else:
                                logger.debug(f"Step {i}: Skipping ECE calculation (no valid data in window).")
                                
                    except Exception as ece_err:
                         logger.warning(f"Step {i}: Could not calculate ECE for coverage alarm: {ece_err}")
                    # --- End ECE Check --- #
                # Coverage Alarm Check
                if self.coverage_alarm_enabled:
                    # Is prediction edge >= threshold?
                    has_edge = abs(edge_step) >= fold_edge_threshold 
                    # Was the prediction correct? (Handle ternary/binary)
                    pred_dir = np.sign(edge_step) if abs(edge_step) >= 1e-6 else 0
                    true_dir = np.sign(step_actual_return) if not is_ternary else np.argmax(true_label_step)-1 # Convert one-hot to -1,0,1
                    is_correct = (pred_dir == true_dir)
                    
                    if has_edge:
                         coverage_hits.append(is_correct)
                         # Trim buffer
                         if len(coverage_hits) > self.coverage_alarm_window:
                              coverage_hits.pop(0)
                         # Check alarm condition
                         if len(coverage_hits) == self.coverage_alarm_window:
                              current_hit_rate = np.mean(coverage_hits)
                              alarm_trigger_level = fold_edge_threshold - self.coverage_alarm_threshold_drop
                              # Note: This condition might need refinement. Comparing hit rate to edge threshold directly isn't ideal.
                              # A better approach might compare current hit rate to expected (e.g., fold validation hit rate)
                              # For now, using the simpler logic from description:
                              if current_hit_rate < (0.5 + alarm_trigger_level/2.0): # Simplified check: hit rate < expected rate at threshold
                                   logger.warning(f"Fold {fold_num} Step {i+1}: Coverage Alarm Triggered! Hit rate {current_hit_rate:.3f} < {0.5 + alarm_trigger_level/2.0:.3f} (Threshold {fold_edge_threshold:.3f}, Drop {self.coverage_alarm_threshold_drop:.3f}) over last {self.coverage_alarm_window} edge samples. Forcing recalibration.")
                                   recalibrate_now = True
                                   coverage_hits = [] # Reset alarm buffer

                # Interval Check
                if not recalibrate_now and (i + 1) % self.recalibrate_every_n == 0:
                    recalibrate_now = True

                # Perform Recalibration
                if recalibrate_now and len(historical_data) >= self.recalibration_window:
                    logger.info(f"Fold {fold_num} Step {i+1}: Recalibrating...")
                    
                    # Get recent data
                    recent_data = historical_data[-self.recalibration_window:]
                    recent_preds, recent_labels = zip(*recent_data)
                    recent_preds = np.array(recent_preds)
                    recent_labels = np.array(recent_labels)

                    if is_ternary and vector_calibrator is not None:
                         vector_calibrator.fit(recent_preds, recent_labels) # Refit
                         if np.array_equal(current_vector_params, vector_calibrator.optimal_params):
                             logger.info("  Vector params unchanged.")
                         else:
                             current_vector_params = vector_calibrator.optimal_params
                             logger.info(f"  New vector params obtained (shape: {current_vector_params.shape}).")
                             # TODO: Optionally save updated params?
                    elif not is_ternary and calibrator is not None:
                         new_T = calibrator.optimise_temperature(recent_preds, recent_labels)
                         if new_T is not None and abs(new_T - current_optimal_T) > 1e-4:
                              logger.info(f"  New optimal temperature: {new_T:.4f} (Previous: {current_optimal_T:.4f})")
                              current_optimal_T = new_T
                              calibrator.optimal_T = new_T # Update instance
                    else:
                              logger.info(f"  Temperature unchanged or optimization failed (T={current_optimal_T:.4f}).")
                    
                    last_recalib_step = i
                elif recalibrate_now:
                     logger.warning(f"Fold {fold_num} Step {i+1}: Recalibration triggered but not enough historical data ({len(historical_data)} < {self.recalibration_window}).")


            # --- Step i: Check for Ruin ---
            if capital <= 0:
                logger.warning(f"Fold {fold_num}: Capital depleted at step {i+1}. Stopping backtest.")
                n_test = i + 1 # Adjust length to current step
                break
                
            # --- Step i: Log Periodic Metrics ---
            if (i + 1) % 1000 == 0: # Log every 1000 steps
                hr_nonzero = (step_correct_nonzero / step_count_nonzero) if step_count_nonzero > 0 else 0
                mean_abs_action = np.mean(step_abs_actions) if step_abs_actions else 0
                metrics_log.append({
                    'step': i + 1, 'timestamp': test_indices[i],
                    'hit_rate_nonzero': hr_nonzero, 'mean_abs_action': mean_abs_action,
                    'equity': capital, 'current_T': current_optimal_T,
                    'last_recalib_step': last_recalib_step
                })
                # Reset interval counters
                step_correct_nonzero, step_count_nonzero, step_abs_actions = 0, 0, []
            # --- End Periodic Log ---
                
        logger.info(f"Fold {fold_num}: Backtest simulation loop finished.")
        logger.info(f"Fold {fold_num}: Final Equity: {capital:.2f}")

        # --- Convert metrics log to DataFrame ---
        metrics_log_df = pd.DataFrame(metrics_log).set_index('step')

        # 5. Prepare Results DataFrame
        if n_test == 0:
             logger.warning(f"Fold {fold_num}: Backtest executed 0 steps.")
             return pd.DataFrame(), {}, pd.DataFrame()
             
        results_data = {
            'equity': equity_curve[1:], 
            'position': positions[1:], 
            'action': actions_taken[1:], 
            'pnl': pnl_steps,           
            'actual_return': actual_ret_test[:n_test],
            'mu_pred': mu_test_state[:n_test],       # Use state mu
            'sigma_pred': sigma_test_state[:n_test], # Use state sigma
            'p_cal_pred': calibrated_probs_steps[:n_test], # Use step-calibrated probs
            'edge_pred': edge_steps[:n_test],           # Use step-calculated edge
            'actual_dir': actual_dir_test[:n_test] if not is_ternary else np.argmax(actual_dir_test[:n_test], axis=1) # Store labels
        }
        if original_prices is not None:
            aligned_prices = original_prices.reindex(test_indices[:n_test]) # Reindex safely
            results_data['close_price'] = aligned_prices['close'].values # Add close price
            
        self.results_df = pd.DataFrame(results_data, index=test_indices[:n_test])
        self.results_df['returns'] = self.results_df['equity'].pct_change().fillna(0.0)
        self.results_df['cumulative_return'] = (1 + self.results_df['returns']).cumprod() - 1
        
        # Calculate Buy & Hold Benchmark
        bh_sharpe = 0.0
        if 'close_price' in self.results_df.columns and not self.results_df['close_price'].isnull().all():
            bh_returns = self.results_df['close_price'].pct_change().fillna(0.0)
            self.results_df['bh_cumulative_return'] = (1 + bh_returns).cumprod() - 1
            bh_sharpe = calculate_sharpe_ratio(bh_returns)
        else:
            self.results_df['bh_cumulative_return'] = 0.0
            logger.warning(f"Fold {fold_num}: Could not calculate Buy & Hold benchmark due to missing/NaN price data.")

        # 6. Calculate Final Performance Metrics
        logger.info(f"Fold {fold_num}: Calculating final performance metrics...")
        final_equity = self.results_df['equity'].iloc[-1]
        total_return_pct = (final_equity / self.initial_capital - 1) * 100
        sharpe = calculate_sharpe_ratio(self.results_df['returns'])
        max_dd = calculate_max_drawdown(self.results_df['equity'])
        
        wins = self.results_df[self.results_df['pnl'] > 0]['pnl']
        losses = self.results_df[self.results_df['pnl'] < 0]['pnl']
        profit_factor = wins.sum() / abs(losses.sum()) if losses.sum() != 0 else np.inf
        
        num_trades = len(trades_executed)
        win_rate_steps = (self.results_df['pnl'] > 0).mean() * 100 if len(self.results_df) > 0 else 0

        # Use the determined fold edge threshold for final metric calculation
        edge_acc, edge_n = edge_filtered_accuracy(
            y_true=self.results_df['actual_dir'], 
            p_cal=self.results_df['p_cal_pred'], 
            thr=fold_edge_threshold 
        )
        sharpe_recentered = calculate_sharpe_ratio(self.results_df['returns'])
        
        # Brier score (only for binary)
        brier = np.nan
        if not is_ternary:
            brier = calculate_brier_score(self.results_df['actual_dir'], self.results_df['p_cal_pred'])

        self.metrics = {
            "Fold Number": fold_num,
            "Test Period Start": test_indices[0].strftime('%Y-%m-%d %H:%M'),
            "Test Period End": test_indices[n_test-1].strftime('%Y-%m-%d %H:%M'),
            "Initial Capital": self.initial_capital,
            "Final Equity": final_equity,
            "Total Net PnL": self.results_df['pnl'].sum(),
            "Total Return (%)": total_return_pct,
            "Annualized Sharpe Ratio": sharpe,
            "Annualized Sharpe Ratio (Re-centred)": sharpe_recentered,
            "Max Drawdown (%)": max_dd * 100,
            "Profit Factor": profit_factor,
            "Number of Trades": num_trades,
            "Win Rate (%)": win_rate_steps, # Use step PnL win rate
            "Edge Threshold Used": fold_edge_threshold,
            "Edge Filtered Accuracy": edge_acc,
            "Edge Filtered N": edge_n,
            "Brier Score": brier,
            "Buy & Hold Sharpe Ratio": bh_sharpe,
            "Initial Optimal T": initial_optimal_T,
            # Add rolling cal info if enabled
            "Rolling Calibration Enabled": self.rolling_cal_enabled,
            "Last Recalibration Step": last_recalib_step if self.rolling_cal_enabled else "N/A",
        }
        logger.info(f"Fold {fold_num}: --- Backtest Simulation Finished ---")
        return self.results_df, self.metrics, metrics_log_df

    def save_results(
        self,
        results_df: pd.DataFrame,
        metrics: Dict[str, Any],
        results_dir: str, # Base results dir for the fold or run
        run_id: str,      # Overall pipeline run_id
        metrics_log_df: Optional[pd.DataFrame] = None,
        fold_num: Optional[int] = None # Pass fold number for unique filenames
    ):
        """
        Saves the backtest results, metrics report, and plots for a specific fold.
        """
        fold_suffix = f"_fold_{fold_num}" if fold_num is not None else ""
        logger.info(f"--- Fold {fold_num}: Saving Backtest Results --- ")
        if results_df is None or metrics is None:
             logger.warning(f"Fold {fold_num}: No results DataFrame or metrics to save.")
             return
             
        if not self.io:
            logger.error(f"Fold {fold_num}: IOManager not provided. Cannot save results.")
            return 

        # Define base filenames with fold suffix
        metrics_fname = f"performance_metrics{fold_suffix}"
        results_df_fname = f"backtest_results{fold_suffix}"
        metrics_log_fname = f"backtest_metrics_log{fold_suffix}"
        summary_plot_fname = f"backtest_summary{fold_suffix}"
        
        # IOManager handles the full path construction including run_id and section
        # We save these fold-specific results within the main run's 'results' section

        # 1. Save Metrics Report 
        try:
            self.io.save_json(metrics, metrics_fname, section='results', use_txt=True)
            logger.info(f"Fold {fold_num}: Performance metrics saved to {self.io.path('results', metrics_fname, suffix='.txt')}")
        except Exception as e:
            logger.error(f"Fold {fold_num}: Failed to save metrics report: {e}", exc_info=True)

        # 2. Save Results DataFrame
        try:
            # Let IOManager decide csv/parquet based on its internal logic
            saved_path = self.io.save_df(results_df, results_df_fname, section='results')
            logger.info(f"Fold {fold_num}: Detailed backtest results saved to {saved_path}")
        except Exception as e:
            logger.error(f"Fold {fold_num}: Failed to save results DataFrame: {e}", exc_info=True)
        
        # 3. Save Metrics Log DataFrame
        if metrics_log_df is not None and not metrics_log_df.empty:
            try:
                saved_path = self.io.save_df(metrics_log_df, metrics_log_fname, section='results')
                logger.info(f"Fold {fold_num}: Periodic backtest metrics log saved to {saved_path}")
            except Exception as e:
                 logger.error(f"Fold {fold_num}: Failed to save metrics log DataFrame: {e}", exc_info=True)
        
        # 4. Generate and Save Plots
        if self.config.get('control', {}).get('generate_plots', True):
            logger.info(f"Fold {fold_num}: Generating backtest plots...")
            try:
                # Plot 1: Multi-subplot summary
                fig_size = self.config.get('output', {}).get('figure_size', [16, 9])
                fig, axes = plt.subplots(3, 1, figsize=fig_size, sharex=True)
                plt.style.use('seaborn-v0_8-darkgrid')
                footer_text = f"© GRU-SAC v3 | Run: {run_id} | Fold: {fold_num}"
                
                # --- Pane 1: Price + Edge Background --- #
                ax = axes[0]
                if 'close_price' in results_df.columns and not results_df['close_price'].isnull().all():
                    ax.plot(results_df.index, results_df['close_price'], label='Price', color='black', alpha=0.9, linewidth=1.0)
                    ax.set_ylabel("Price")
                    
                    edge_thr = metrics.get("Edge Threshold Used", 0.1) # Get actual threshold used
                    long_edge_mask = results_df['edge_pred'] >= edge_thr
                    short_edge_mask = results_df['edge_pred'] <= -edge_thr
                    ax.fill_between(results_df.index, ax.get_ylim()[0], ax.get_ylim()[1], 
                                    where=long_edge_mask, color='blue', alpha=0.1, label=f'Long Edge >= {edge_thr:.2f}')
                    ax.fill_between(results_df.index, ax.get_ylim()[0], ax.get_ylim()[1], 
                                    where=short_edge_mask, color='red', alpha=0.1, label=f'Short Edge <= {-edge_thr:.2f}')
                else:
                     ax.text(0.5, 0.5, 'Price data unavailable', ha='center', va='center')
                ax.set_title(f'Backtest Summary (Fold: {fold_num})', fontsize=14)
                ax.legend(fontsize=8)
                ax.grid(True, linestyle='--', alpha=0.6)

                # --- Pane 2: Position Size --- #
                ax = axes[1]
                ax.plot(results_df.index, results_df['position'], label='Target Position', color='purple', drawstyle='steps-post')
                ax.set_ylabel("Position (-1 to 1)")
                ax.set_ylim(-1.1, 1.1)
                ax.legend(fontsize=8)
                ax.grid(True, linestyle='--', alpha=0.6)

                # --- Pane 3: Equity Curve + Drawdowns --- #
                ax = axes[2]
                equity_norm = results_df['equity'] / self.initial_capital
                ax.plot(results_df.index, equity_norm, label='Strategy Equity', color='green')
                # Add Buy & Hold if available
                if 'bh_cumulative_return' in results_df.columns:
                     ax.plot(results_df.index, results_df['bh_cumulative_return'] + 1, label='Buy & Hold Equity', color='grey', linestyle=':')
                ax.set_ylabel("Normalized Equity")
                ax.set_xlabel("Time")
                
                rolling_max_norm = equity_norm.cummax()
                drawdown_norm = (equity_norm - rolling_max_norm) 
                ax.fill_between(results_df.index, equity_norm, rolling_max_norm, where=drawdown_norm < 0, 
                                color='red', alpha=0.3, label='Drawdown')
                
                sharpe_val = metrics.get('Annualized Sharpe Ratio (Re-centred)', metrics.get('Annualized Sharpe Ratio', np.nan))
                max_dd_val = metrics.get('Max Drawdown (%)', np.nan)
                metrics_text = f"Sharpe: {sharpe_val:.2f}\nMax DD: {max_dd_val:.2f}%"
                ax.text(0.02, 0.1, metrics_text, transform=ax.transAxes, fontsize=9,
                        verticalalignment='bottom', bbox=dict(boxstyle='round,pad=0.5', fc='wheat', alpha=0.5))
                
                ax.legend(fontsize=8)
                ax.grid(True, linestyle='--', alpha=0.6)
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M'))
                plt.xticks(rotation=30)

                fig.text(0.99, 0.01, footer_text, horizontalalignment='right', 
                         verticalalignment='bottom', fontsize=8, color='gray')
                
                plt.tight_layout(rect=[0, 0.03, 1, 0.95]) 
                
                # Save using IOManager
                saved_path = self.io.save_figure(fig, summary_plot_fname, section='results')
                logger.info(f"Fold {fold_num}: Backtest summary plot saved to {saved_path}")
                plt.close(fig)

            except Exception as e:
                 logger.error(f"Fold {fold_num}: Failed to generate or save backtest plots: {e}", exc_info=True)
        else:
            logger.info(f"Fold {fold_num}: Skipping plot generation as per config.")
            
        logger.info(f"--- Fold {fold_num}: Finished Saving Backtest Results ---") 

    # --- Helper for ECE Calculation (Revision 4) --- #
    def _calculate_ece(self, probs: np.ndarray, y_true_onehot: np.ndarray, n_bins: int = 10) -> float:
        """Calculates the Expected Calibration Error (ECE) for multi-class models.

        Args:
            probs (np.ndarray): Predicted probabilities (N x K).
            y_true_onehot (np.ndarray): True labels, one-hot encoded (N x K).
            n_bins (int): Number of bins to divide the confidence scores into.

        Returns:
            float: The calculated ECE.
        """
        if probs.shape != y_true_onehot.shape:
            raise ValueError("Probs and y_true_onehot must have the same shape.")
        if len(probs.shape) != 2:
            raise ValueError("Inputs must be 2D arrays (N x K).")
            
        num_samples = probs.shape[0]
        confidences = np.max(probs, axis=1)
        predictions = np.argmax(probs, axis=1)
        true_labels = np.argmax(y_true_onehot, axis=1)
        accuracies = (predictions == true_labels).astype(float)

        ece = 0.0
        bin_lowers = np.linspace(0.0, 1.0, n_bins + 1)[:-1]
        bin_uppers = np.linspace(0.0, 1.0, n_bins + 1)[1:]

        for bin_lower, bin_upper in zip(bin_lowers, bin_uppers):
            in_bin = (confidences > bin_lower) & (confidences <= bin_upper)
            prop_in_bin = np.mean(in_bin)

            if prop_in_bin > 0:
                accuracy_in_bin = np.mean(accuracies[in_bin])
                avg_confidence_in_bin = np.mean(confidences[in_bin])
                ece += np.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
                
        return ece
    # --- End ECE Helper --- # 