"""
Contains the BaselineChecker class for running baseline model checks.
"""

import logging
import pandas as pd
import numpy as np
import scipy.stats as st
from scipy.stats import binomtest
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier # Import RF
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
from typing import Dict, Any, Optional, Tuple, Union, List
import matplotlib.pyplot as plt
import seaborn as sns
import os
# Add CalibratedClassifierCV import
from sklearn.calibration import CalibratedClassifierCV 

# Import IOManager for saving results
try:
    from ..io_manager import IOManager
except ImportError:
    logging.warning("IOManager could not be imported in BaselineChecker. Report saving will fail.")
    IOManager = None # Placeholder if unavailable

logger = logging.getLogger(__name__)

class BaselineChecker:
    """Runs baseline model checks, including raw and edge-filtered evaluations."""

    def __init__(self, config: Dict[str, Any], io: IOManager):
        """
        Initialize the BaselineChecker.
        
        Args:
            config (Dict[str, Any]): Pipeline configuration dictionary.
            io: IOManager instance for saving.
        """
        self.config = config
        self.io = io

    def _calculate_ci_lower_bound(self, k: int, n: int, p: float, conf_level: float) -> float:
        """Helper to calculate binomial test CI lower bound."""
        if n == 0: return np.nan
        try:
            return binomtest(k, n, p=p, alternative='greater').proportion_ci(confidence_level=conf_level).low
        except ValueError as e:
            logger.error(f"Error calculating binomial test (k={k}, n={n}, p={p}): {e}. Returning NaN.")
            return np.nan

    def run_binary_logistic_baseline(self, 
                                     X_train_scaled: pd.DataFrame, 
                                     y_train_dir_binary: pd.Series, 
                                     X_val_scaled: pd.DataFrame, 
                                     y_val_dir_binary: pd.Series,
                                     io: Optional[IOManager] = None
                                     ) -> Tuple[Optional[LogisticRegression], Any, Dict[str, Any]]: # Return base_model, calibrator, report
        """
        Runs Binary Logistic Regression baseline. Uses internal split for raw CI.
        Optionally applies isotonic calibration before edge checks.

        Args:
            X_train_scaled: SCALED training features.
            y_train_dir_binary: BINARY {0, 1} training labels.
            X_val_scaled: SCALED validation features.
            y_val_dir_binary: BINARY {0, 1} validation labels.
            io: IOManager instance for saving.

        Returns:
            Tuple: (base_lr_model, final_predictor, report_dict)
                   base_lr_model: Raw Logistic Regression model.
                   final_predictor: Calibrated model if enabled, else base_lr_model.
                   report_dict: Dictionary containing metrics.
        """
        # --- Configuration --- #
        baseline_cfg = self.config.get('baselines', {}).get('logistic_regression', {})
        max_iter = baseline_cfg.get('max_iter', 1000)
        solver = baseline_cfg.get('solver', 'lbfgs')
        # random_state = baseline_cfg.get('random_state', 42) # Use calibration random state
        val_subset_split_ratio = baseline_cfg.get('val_subset_split_ratio', 0.2)
        val_subset_shuffle = baseline_cfg.get('val_subset_shuffle', False)
        ci_confidence_level = baseline_cfg.get('ci_confidence_level', 0.95)
        # Calibration Config
        calibration_enabled = baseline_cfg.get('calibration_enabled', False)
        calibration_method = baseline_cfg.get('calibration_method', 'isotonic')
        calibration_holdout = baseline_cfg.get('calibration_holdout', 0.2)
        calibration_random_state = baseline_cfg.get('random_state', 42) # Reuse random_state
        
        logger.info(f"Running Binary Logistic Regression baseline (Baseline 1) on {X_train_scaled.shape[1]} scaled features.")
        if calibration_enabled:
             logger.info(f"  Calibration enabled: method={calibration_method}, holdout={calibration_holdout}")
        else:
             logger.info("  Calibration disabled.")

        # --- Handle potential NaNs and Infs in input --- #
        # ... (NaN/Inf handling code remains the same) ...
        logger.debug(f"Checking for NaNs/Infs before LogReg. X_train_scaled NaNs: {X_train_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_train_scaled).sum().sum()}")
        X_train_scaled_clean = pd.DataFrame(np.nan_to_num(X_train_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_train_scaled.index, columns=X_train_scaled.columns)
        logger.debug(f"Checking for NaNs/Infs before LogReg. X_val_scaled NaNs: {X_val_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_val_scaled).sum().sum()}")
        X_val_scaled_clean = pd.DataFrame(np.nan_to_num(X_val_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_val_scaled.index, columns=X_val_scaled.columns)
        # Add check after cleaning
        if X_train_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_train_scaled_clean).sum().sum() > 0:
            logger.warning("NaNs or Infs persisted after np.nan_to_num in run_binary_logistic_baseline (Train).")
        if X_val_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_val_scaled_clean).sum().sum() > 0:
             logger.warning("NaNs or Infs persisted after np.nan_to_num in run_binary_logistic_baseline (Val).")
        # --- End Handle NaNs/Infs ---

        # --- Initialize Report --- #
        report = {
            "baseline_model_type": "LogisticRegression_Binary",
            "calibration_enabled": calibration_enabled,
            "accuracy_val_subset": np.nan,
            "ci_lower_bound": np.nan, 
            "n_val_subset": 0,
            "accuracy_orig_val": np.nan,
            "classification_report_orig_val": "N/A",
        }
        
        base_lr_model = None 
        calibrator = None
        final_predictor = None # This will hold the calibrated model if enabled, else the base model

        try:
            # --- Split train data further if calibration is enabled --- #
            if calibration_enabled:
                 logger.info(f"Splitting training data for calibration (holdout={calibration_holdout})...")
                 X_train_for_lr, X_cal, y_train_for_lr, y_cal = train_test_split(
                     X_train_scaled_clean, y_train_dir_binary,
                     test_size=calibration_holdout,
                     shuffle=True, # Shuffle is recommended for calibration split
                     stratify=y_train_dir_binary,
                     random_state=calibration_random_state
                 )
                 logger.info(f"  Base LR training set size: {X_train_for_lr.shape[0]}, Calibration set size: {X_cal.shape[0]}")
            else:
                 # Use all provided training data for LR
                 X_train_for_lr = X_train_scaled_clean
                 y_train_for_lr = y_train_dir_binary
                 X_cal, y_cal = None, None # No calibration set
            
            # --- Train Base Logistic Regression Model --- #
            logger.info(f"Fitting Base Binary LogReg on {X_train_for_lr.shape[0]} samples...")
            base_lr_model = LogisticRegression(
                 max_iter=max_iter, 
                 solver=solver, 
                 random_state=calibration_random_state # Use same random state
            )
            base_lr_model.fit(X_train_for_lr, y_train_for_lr)
            final_predictor = base_lr_model # Default predictor is the base model
            
            # --- Fit Calibrator (if enabled) --- #
            if calibration_enabled and X_cal is not None and not X_cal.empty:
                 logger.info(f"Fitting Calibrator ({calibration_method}) on {X_cal.shape[0]} samples...")
                 try:
                      calibrator = CalibratedClassifierCV(
                          base_lr_model, 
                          method=calibration_method, 
                          cv='prefit' # Use the already fitted base model
                      )
                      calibrator.fit(X_cal, y_cal)
                      final_predictor = calibrator # Update final predictor to the calibrated one
                      logger.info("Calibration fitting complete.")
                 except Exception as cal_err:
                      logger.error(f"Error fitting calibrator: {cal_err}. Using uncalibrated model.", exc_info=True)
                      final_predictor = base_lr_model # Fallback to base model on calibration error
            elif calibration_enabled:
                 logger.warning("Calibration enabled but calibration set is empty. Using uncalibrated model.")
                 final_predictor = base_lr_model

            # --- Internal Split for Raw CI (Always uses BASE model) --- # 
            logger.info("Performing internal split for Raw CI calculation...")
            # Use the same X_train_for_lr data that the base model was trained on
            X_teach, X_val_subset, y_teach, y_val_subset = train_test_split(
                X_train_for_lr, y_train_for_lr, 
                test_size=val_subset_split_ratio, # Use the original val_subset ratio
                shuffle=val_subset_shuffle,
                stratify=y_train_for_lr if val_subset_shuffle else None,
                random_state=calibration_random_state + 1 # Use different seed for this split
            )
            report["n_val_subset"] = len(y_val_subset)

            if report["n_val_subset"] > 0:
                 # Fit a TEMPORARY base model on the teaching subset for raw CI
                 logger.info(f"Fitting TEMP Binary LogReg on teaching subset ({len(y_teach)} samples) for Raw CI...")
                 temp_lr_model = LogisticRegression(
                     max_iter=max_iter, 
                     solver=solver, 
                     random_state=calibration_random_state # Re-use state
                 )
                 temp_lr_model.fit(X_teach, y_teach) 

                 logger.debug("Predicting on internal validation subset (Binary LogReg - RAW CI)...")
                 # Use predict (not predict_proba) on the TEMP model
                 y_pred_val_subset = temp_lr_model.predict(X_val_subset) 
                 acc_val_subset = accuracy_score(y_val_subset, y_pred_val_subset)
                 report["accuracy_val_subset"] = float(acc_val_subset)
                 
                 # Calculate Raw CI lower bound (vs 0.5) using the TEMP model results
                 k_correct = np.sum(y_pred_val_subset == y_val_subset)
                 n = report["n_val_subset"]
                 raw_ci_lb = self._calculate_ci_lower_bound(k_correct, n, p=0.5, conf_level=ci_confidence_level)
                 report["ci_lower_bound"] = raw_ci_lb
                 logger.info(f"Baseline 1 (Binary LogReg - Raw Internal Split): Acc={acc_val_subset:.3f}, CI_LB={raw_ci_lb:.3f}")
            else:
                logger.warning("Internal validation subset empty for Binary LogReg Raw CI. Skipping calculation.")

            # --- Evaluate on Original Validation Set (Uses FINAL predictor) --- #
            if final_predictor is not None and not X_val_scaled_clean.empty:
                logger.info("Evaluating FINAL predictor (calibrated if enabled) on original validation set...")
                # Use predict for accuracy/report, predict_proba will be used by edge check
                orig_val_pred = final_predictor.predict(X_val_scaled_clean) 
                orig_val_acc = accuracy_score(y_val_dir_binary, orig_val_pred)
                report["accuracy_orig_val"] = float(orig_val_acc)
                try:
                    report["classification_report_orig_val"] = classification_report(y_val_dir_binary, orig_val_pred, output_dict=True, zero_division=0)
                except Exception as report_err:
                     logger.warning(f"Could not generate Binary LogReg classification report: {report_err}")
                     report["classification_report_orig_val"] = "Error generating report"
                logger.info(f"Baseline 1 (Binary LogReg - FINAL Predictor): Original validation set accuracy: {orig_val_acc:.3f}")
            elif final_predictor is None:
                 logger.warning("FINAL predictor (calibrated or base LR) not available. Skipping evaluation on original validation set.")
            else: # X_val_scaled is empty
                 logger.warning("Original validation set is empty. Skipping Binary LogReg evaluation on it.")

        except Exception as e:
            logger.error(f"Failed during Binary Logistic Regression baseline: {e}", exc_info=True)
            report["notes"] = f"Failed with error: {e}"
            base_lr_model = None # Ensure models are None on failure
            final_predictor = None
            
        # --- Save Report --- #
        # ... (Saving logic remains the same) ...
        if io:
            try:
                io.save_json(report, "baseline_binary_lr_report", section='results', use_txt=True)
                logger.info("Saved Binary Logistic Regression baseline report.")
            except Exception as save_e:
                 logger.error(f"Failed to save Binary LogReg baseline report: {save_e}")
        else:
             logger.warning("IOManager not provided, skipping saving Binary LogReg baseline report.")

        # Return the base model (for raw CI check logic if separated) and the final predictor (for edge checks)
        return base_lr_model, final_predictor, report

    def run_binary_random_forest_baseline(self, 
                                          X_train_scaled: pd.DataFrame, 
                                          y_train_dir_binary: pd.Series, 
                                          X_val_scaled: pd.DataFrame, 
                                          y_val_dir_binary: pd.Series,
                                          io: Optional[IOManager] = None
                                          ) -> Tuple[Optional[RandomForestClassifier], Dict[str, Any]]:
        """
        Runs Binary RandomForest baseline. Uses internal split for raw CI.

        Args:
            X_train_scaled: SCALED training features.
            y_train_dir_binary: BINARY {0, 1} training labels.
            X_val_scaled: SCALED validation features.
            y_val_dir_binary: BINARY {0, 1} validation labels.
            io: IOManager instance for saving.

        Returns:
            Tuple: (fitted_model, report_dict)
        """
        # --- Configuration --- #
        rf_cfg = self.config.get('baselines', {}).get('random_forest', {}) # Reuse binary RF config section
        n_estimators = rf_cfg.get('n_estimators', 100)
        max_depth = rf_cfg.get('max_depth', None)
        min_samples_split = rf_cfg.get('min_samples_split', 2)
        min_samples_leaf = rf_cfg.get('min_samples_leaf', 1)
        random_state = rf_cfg.get('random_state', 42)
        n_jobs = rf_cfg.get('n_jobs', -1)
        val_subset_split_ratio = rf_cfg.get('val_subset_split_ratio', 0.2)
        val_subset_shuffle = rf_cfg.get('val_subset_shuffle', False)
        ci_confidence_level = rf_cfg.get('ci_confidence_level', 0.95)

        logger.info(f"Running Binary RandomForest baseline (Baseline 2) on {X_train_scaled.shape[1]} scaled features.")
        logger.info(f"  RF Params: n_estimators={n_estimators}, max_depth={max_depth}, n_jobs={n_jobs}")

        # --- Handle potential NaNs and Infs in input ---
        logger.debug(f"Checking for NaNs/Infs before RF. X_train_scaled NaNs: {X_train_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_train_scaled).sum().sum()}")
        X_train_scaled_clean = pd.DataFrame(np.nan_to_num(X_train_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_train_scaled.index, columns=X_train_scaled.columns)
        logger.debug(f"Checking for NaNs/Infs before RF. X_val_scaled NaNs: {X_val_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_val_scaled).sum().sum()}")
        X_val_scaled_clean = pd.DataFrame(np.nan_to_num(X_val_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_val_scaled.index, columns=X_val_scaled.columns)
        # Add check after cleaning
        if X_train_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_train_scaled_clean).sum().sum() > 0:
            logger.warning("NaNs or Infs persisted after np.nan_to_num in run_binary_random_forest_baseline (Train).")
        if X_val_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_val_scaled_clean).sum().sum() > 0:
             logger.warning("NaNs or Infs persisted after np.nan_to_num in run_binary_random_forest_baseline (Val).")
        # --- End Handle NaNs/Infs ---

        # --- Initialize Report --- #
        report = {
            "baseline_model_type": "RandomForestClassifier_Binary", # Baseline 2 ID
            "accuracy_val_subset": np.nan,
            "ci_lower_bound": np.nan, # Raw CI LB (internal split)
            "n_val_subset": 0,
            "accuracy_orig_val": np.nan,
            "classification_report_orig_val": "N/A",
        }
        
        baseline_model = None
        acc_val_subset = np.nan # Initialize outside the if block

        try:
            # Split train data into teach/validation subsets for raw CI calculation
            X_teach, X_val_subset, y_teach, y_val_subset = train_test_split(
                X_train_scaled_clean, y_train_dir_binary, # Use cleaned data
                test_size=val_subset_split_ratio, 
                shuffle=val_subset_shuffle,
                stratify=y_train_dir_binary if val_subset_shuffle else None
            )
            report["n_val_subset"] = len(y_val_subset)

            if report["n_val_subset"] > 0:
                logger.info(f"Fitting Binary RF on teaching subset ({len(y_teach)} samples)...")
                baseline_model = RandomForestClassifier(
                    n_estimators=n_estimators,
                    max_depth=max_depth,
                    min_samples_split=min_samples_split,
                    min_samples_leaf=min_samples_leaf,
                    random_state=random_state,
                    n_jobs=n_jobs
                )
                baseline_model.fit(X_teach, y_teach) # Fit on cleaned subset
            
                logger.debug("Predicting on internal validation subset (Binary RF)..")
                y_pred_val_subset = baseline_model.predict(X_val_subset) # Predict on cleaned subset
                acc_val_subset = accuracy_score(y_val_subset, y_pred_val_subset)
                report["accuracy_val_subset"] = float(acc_val_subset) # Now correctly indented inside the if
            
                # Calculate Raw CI lower bound (vs 0.5) - Correctly indented inside the if
                k_correct = np.sum(y_pred_val_subset == y_val_subset)
                n = report["n_val_subset"]
                raw_ci_lb = self._calculate_ci_lower_bound(k_correct, n, p=0.5, conf_level=ci_confidence_level)
                report["ci_lower_bound"] = raw_ci_lb
                logger.info(f"Baseline 2 (Binary RF - Raw Internal Split): Acc={acc_val_subset:.3f}, CI_LB={raw_ci_lb:.3f}")
            else: # Correctly aligned with if
                # Correctly indented inside the else
                logger.warning("Internal validation subset empty for Binary RF. Skipping raw CI calculation.")

            # --- Evaluate on Original Validation Set --- #
            if baseline_model is not None and not X_val_scaled_clean.empty:
                logger.info("Evaluating Binary RF on original validation set...")
                orig_val_pred = baseline_model.predict(X_val_scaled_clean) # Predict on cleaned validation data
                orig_val_acc = accuracy_score(y_val_dir_binary, orig_val_pred)
                report["accuracy_orig_val"] = float(orig_val_acc)
                try:
                    report["classification_report_orig_val"] = classification_report(y_val_dir_binary, orig_val_pred, output_dict=True, zero_division=0)
                except Exception as report_err:
                     logger.warning(f"Could not generate Binary RF classification report: {report_err}")
                     report["classification_report_orig_val"] = "Error generating report"
                logger.info(f"Baseline 2 (Binary RF): Original validation set accuracy: {orig_val_acc:.3f}")
            elif baseline_model is None:
                 logger.warning("Binary RF model not trained. Skipping evaluation on original validation set.")
            else: # X_val_scaled is empty
                 logger.warning("Original validation set is empty. Skipping Binary RF evaluation on it.")

        except ImportError:
            logger.error("Scikit-learn (RandomForestClassifier) not found. Cannot run Baseline 2.")
            report["notes"] = "Skipped - scikit-learn not installed."
            baseline_model = None
        except Exception as e:
            logger.error(f"Failed during Binary RandomForest baseline: {e}", exc_info=True)
            report["notes"] = f"Failed with error: {e}"
            baseline_model = None
        
        # --- Save RF Baseline Report --- #
        if io:
            try:
                io.save_json(report, "baseline_binary_rf_report", section='results', use_txt=True)
                logger.info("Saved Binary RandomForest baseline report.")
            except Exception as save_e:
                 logger.error(f"Failed to save Binary RF baseline report: {save_e}")
        else:
             logger.warning("IOManager not provided, skipping saving Binary RF baseline report.")
        
        return baseline_model, report

    def run_ternary_logistic_baseline(self,
                                      X_train_scaled: pd.DataFrame,
                                      y_train_dir_ordinal: pd.Series,
                                      X_val_scaled: pd.DataFrame,
                                      y_val_dir_ordinal: pd.Series,
                                      io: Optional[IOManager] = None
                                      ) -> Tuple[Optional[LogisticRegression], Dict[str, Any]]:
        """
        Runs Multinomial Logistic Regression baseline for ternary labels. Uses internal split for raw CI.

        Args:
            X_train_scaled: SCALED training features.
            y_train_dir_ordinal: ORDINAL {0, 1, 2} training labels.
            X_val_scaled: SCALED validation features.
            y_val_dir_ordinal: ORDINAL {0, 1, 2} validation labels.
            io: IOManager instance for saving.

        Returns:
            Tuple: (fitted_model, report_dict)
        """
        # --- Configuration --- #
        mlogreg_cfg = self.config.get('baselines', {}).get('multinomial_logistic_regression', {})
        max_iter = mlogreg_cfg.get('max_iter', 1000)
        solver = mlogreg_cfg.get('solver', 'lbfgs')
        multi_class = mlogreg_cfg.get('multi_class', 'multinomial')
        random_state = mlogreg_cfg.get('random_state', 42)
        val_subset_split_ratio = mlogreg_cfg.get('val_subset_split_ratio', 0.2)
        val_subset_shuffle = mlogreg_cfg.get('val_subset_shuffle', False)
        ci_confidence_level = mlogreg_cfg.get('ci_confidence_level', 0.95)
        chance_level = 1/3 # For ternary CI

        logger.info(f"Running Ternary Logistic Regression baseline (Baseline 3) on {X_train_scaled.shape[1]} scaled features.")

        # --- Handle potential NaNs and Infs in input ---
        logger.debug(f"Checking for NaNs/Infs before Ternary LogReg. X_train_scaled NaNs: {X_train_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_train_scaled).sum().sum()}")
        X_train_scaled_clean = pd.DataFrame(np.nan_to_num(X_train_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_train_scaled.index, columns=X_train_scaled.columns)
        logger.debug(f"Checking for NaNs/Infs before Ternary LogReg. X_val_scaled NaNs: {X_val_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_val_scaled).sum().sum()}")
        X_val_scaled_clean = pd.DataFrame(np.nan_to_num(X_val_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_val_scaled.index, columns=X_val_scaled.columns)
        # Add check after cleaning
        if X_train_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_train_scaled_clean).sum().sum() > 0:
            logger.warning("NaNs or Infs persisted after np.nan_to_num in run_ternary_logistic_baseline (Train).")
        if X_val_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_val_scaled_clean).sum().sum() > 0:
             logger.warning("NaNs or Infs persisted after np.nan_to_num in run_ternary_logistic_baseline (Val).")
        # --- End Handle NaNs/Infs ---

        # --- Initialize Report --- #
        report = {
            "baseline_model_type": "LogisticRegression_Ternary", # Baseline 3 ID
            "accuracy_val_subset": np.nan,
            "ci_lower_bound": np.nan, # Raw CI LB (internal split vs 1/3)
            "n_val_subset": 0,
            "accuracy_orig_val": np.nan,
            "classification_report_orig_val": "N/A",
        }
        
        baseline_model = None

        try:
            # Split train data into teach/validation subsets for raw CI calculation
            X_teach, X_val_subset, y_teach, y_val_subset = train_test_split(
                X_train_scaled_clean, y_train_dir_ordinal, # Use cleaned data
                test_size=val_subset_split_ratio, 
                shuffle=val_subset_shuffle,
                stratify=y_train_dir_ordinal if val_subset_shuffle else None
            )
            report["n_val_subset"] = len(y_val_subset)

            if report["n_val_subset"] > 0:
                logger.info(f"Fitting Ternary LogReg on teaching subset ({len(y_teach)} samples)...")
                baseline_model = LogisticRegression(
                    max_iter=max_iter, 
                    solver=solver, 
                    multi_class=multi_class,
                    random_state=random_state
                )
                baseline_model.fit(X_teach, y_teach) # Fit on cleaned subset
                
                logger.debug("Predicting on internal validation subset (Ternary LogReg)...")
                y_pred_val_subset = baseline_model.predict(X_val_subset) # Predict on cleaned subset
                acc_val_subset = accuracy_score(y_val_subset, y_pred_val_subset)
                report["accuracy_val_subset"] = float(acc_val_subset)
                
                # Calculate Raw CI lower bound (vs 1/3)
                k_correct = np.sum(y_pred_val_subset == y_val_subset)
                n = report["n_val_subset"]
                raw_ci_lb = self._calculate_ci_lower_bound(k_correct, n, p=chance_level, conf_level=ci_confidence_level)
                report["ci_lower_bound"] = raw_ci_lb
                logger.info(f"Baseline 3 (Ternary LogReg - Raw Internal Split): Acc={acc_val_subset:.3f}, CI_LB={raw_ci_lb:.3f} (vs {chance_level:.3f})")
            else:
                logger.warning("Internal validation subset empty for Ternary LogReg. Skipping raw CI calculation.")

            # Evaluate on original validation set
            if baseline_model is not None and not X_val_scaled_clean.empty:
                logger.info("Evaluating Ternary LogReg on original validation set...")
                orig_val_pred = baseline_model.predict(X_val_scaled_clean) # Predict on cleaned validation data
                orig_val_acc = accuracy_score(y_val_dir_ordinal, orig_val_pred)
                report["accuracy_orig_val"] = float(orig_val_acc)
                try:
                    report["classification_report_orig_val"] = classification_report(y_val_dir_ordinal, orig_val_pred, output_dict=True, zero_division=0)
                except Exception as report_err:
                     logger.warning(f"Could not generate Ternary LogReg classification report: {report_err}")
                     report["classification_report_orig_val"] = "Error generating report"
                logger.info(f"Baseline 3 (Ternary LogReg): Original validation set accuracy: {orig_val_acc:.3f}")
            elif baseline_model is None:
                logger.warning("Ternary LogReg model not trained. Skipping evaluation on original validation set.")
            else:
                logger.warning("Original validation set empty for Ternary LogReg.")

        except Exception as e:
            logger.error(f"Failed during Ternary Logistic Regression baseline: {e}", exc_info=True)
            report["notes"] = f"Failed with error: {e}"
            baseline_model = None
            
        # Save Report
        if io:
            try:
                io.save_json(report, "baseline_ternary_lr_report", section='results', use_txt=True)
                logger.info("Saved Ternary Logistic Regression baseline report.")
            except Exception as save_e:
                 logger.error(f"Failed to save Ternary LogReg report: {save_e}")
        else:
             logger.warning("IOManager not provided, skipping saving Ternary LogReg report.")

        return baseline_model, report
        
    def run_ternary_random_forest_baseline(self,
                                           X_train_scaled: pd.DataFrame,
                                           y_train_dir_ordinal: pd.Series,
                                           X_val_scaled: pd.DataFrame,
                                           y_val_dir_ordinal: pd.Series,
                                           io: Optional[IOManager] = None
                                           ) -> Tuple[Optional[RandomForestClassifier], Dict[str, Any]]:
        """
        Runs Ternary RandomForest baseline. Uses internal split for raw CI.

        Args:
            X_train_scaled: SCALED training features.
            y_train_dir_ordinal: ORDINAL {0, 1, 2} training labels.
            X_val_scaled: SCALED validation features.
            y_val_dir_ordinal: ORDINAL {0, 1, 2} validation labels.
            io: IOManager instance for saving.

        Returns:
            Tuple: (fitted_model, report_dict)
        """
        # --- Configuration --- #
        rf_cfg = self.config.get('baselines', {}).get('ternary_random_forest', {})
        n_estimators = rf_cfg.get('n_estimators', 100)
        max_depth = rf_cfg.get('max_depth', None)
        min_samples_split = rf_cfg.get('min_samples_split', 2)
        min_samples_leaf = rf_cfg.get('min_samples_leaf', 1)
        random_state = rf_cfg.get('random_state', 42)
        n_jobs = rf_cfg.get('n_jobs', -1)
        val_subset_split_ratio = rf_cfg.get('val_subset_split_ratio', 0.2)
        val_subset_shuffle = rf_cfg.get('val_subset_shuffle', False)
        ci_confidence_level = rf_cfg.get('ci_confidence_level', 0.95)
        chance_level = 1/3

        logger.info(f"Running Ternary RandomForest baseline (Baseline 4) on {X_train_scaled.shape[1]} scaled features.")
        logger.info(f"  RF Params: n_estimators={n_estimators}, max_depth={max_depth}, n_jobs={n_jobs}")

        # --- Handle potential NaNs and Infs in input ---
        logger.debug(f"Checking for NaNs/Infs before Ternary RF. X_train_scaled NaNs: {X_train_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_train_scaled).sum().sum()}")
        X_train_scaled_clean = pd.DataFrame(np.nan_to_num(X_train_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_train_scaled.index, columns=X_train_scaled.columns)
        logger.debug(f"Checking for NaNs/Infs before Ternary RF. X_val_scaled NaNs: {X_val_scaled.isnull().sum().sum()}, Infs: {np.isinf(X_val_scaled).sum().sum()}")
        X_val_scaled_clean = pd.DataFrame(np.nan_to_num(X_val_scaled, nan=0.0, posinf=0.0, neginf=0.0), index=X_val_scaled.index, columns=X_val_scaled.columns)
        # Add check after cleaning
        if X_train_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_train_scaled_clean).sum().sum() > 0:
            logger.warning("NaNs or Infs persisted after np.nan_to_num in run_ternary_random_forest_baseline (Train).")
        if X_val_scaled_clean.isnull().sum().sum() > 0 or np.isinf(X_val_scaled_clean).sum().sum() > 0:
             logger.warning("NaNs or Infs persisted after np.nan_to_num in run_ternary_random_forest_baseline (Val).")
        # --- End Handle NaNs/Infs ---

        # --- Initialize Report --- #
        report = {
            "baseline_model_type": "RandomForestClassifier_Ternary", # Baseline 4 ID
            "accuracy_val_subset": np.nan,
            "ci_lower_bound": np.nan, # Raw CI LB (internal split vs 1/3)
            "n_val_subset": 0,
            "accuracy_orig_val": np.nan,
            "classification_report_orig_val": "N/A",
        }
        
        baseline_model = None

        try:
            # Split train data into teach/validation subsets for raw CI calculation
            X_teach, X_val_subset, y_teach, y_val_subset = train_test_split(
                X_train_scaled_clean, y_train_dir_ordinal, # Use cleaned data
                test_size=val_subset_split_ratio, 
                shuffle=val_subset_shuffle,
                stratify=y_train_dir_ordinal if val_subset_shuffle else None
            )
            report["n_val_subset"] = len(y_val_subset)

            if report["n_val_subset"] > 0:
                logger.info(f"Fitting Ternary RF on teaching subset ({len(y_teach)} samples)...")
                baseline_model = RandomForestClassifier(
                    n_estimators=n_estimators,
                    max_depth=max_depth,
                    min_samples_split=min_samples_split,
                    min_samples_leaf=min_samples_leaf,
                    random_state=random_state,
                    n_jobs=n_jobs
                )
                baseline_model.fit(X_teach, y_teach) # Fit on cleaned subset
                
                logger.debug("Predicting on internal validation subset (Ternary RF)..")
                y_pred_val_subset = baseline_model.predict(X_val_subset) # Predict on cleaned subset
                acc_val_subset = accuracy_score(y_val_subset, y_pred_val_subset)
                report["accuracy_val_subset"] = float(acc_val_subset)
                
                # Calculate Raw CI lower bound (vs 1/3)
                k_correct = np.sum(y_pred_val_subset == y_val_subset)
                n = report["n_val_subset"]
                raw_ci_lb = self._calculate_ci_lower_bound(k_correct, n, p=chance_level, conf_level=ci_confidence_level)
                report["ci_lower_bound"] = raw_ci_lb
                logger.info(f"Baseline 4 (Ternary RF - Raw Internal Split): Acc={acc_val_subset:.3f}, CI_LB={raw_ci_lb:.3f} (vs {chance_level:.3f})")
            else:
                logger.warning("Internal validation subset empty for Ternary RF. Skipping raw CI calculation.")

            # Evaluate on original validation set
            if baseline_model is not None and not X_val_scaled_clean.empty:
                logger.info("Evaluating Ternary RF on original validation set...")
                orig_val_pred = baseline_model.predict(X_val_scaled_clean) # Predict on cleaned validation data
                orig_val_acc = accuracy_score(y_val_dir_ordinal, orig_val_pred)
                report["accuracy_orig_val"] = float(orig_val_acc)
                try:
                    report["classification_report_orig_val"] = classification_report(y_val_dir_ordinal, orig_val_pred, output_dict=True, zero_division=0)
                except Exception as report_err:
                     logger.warning(f"Could not generate Ternary RF classification report: {report_err}")
                     report["classification_report_orig_val"] = "Error generating report"
                logger.info(f"Baseline 4 (Ternary RF): Original validation set accuracy: {orig_val_acc:.3f}")
            elif baseline_model is None:
                logger.warning("Ternary RF model not trained. Skipping evaluation on original validation set.")
            else:
                logger.warning("Original validation set empty for Ternary RF.")

        except ImportError:
            logger.error("Scikit-learn (RandomForestClassifier) not found. Cannot run Baseline 4.")
            report["notes"] = "Skipped - scikit-learn not installed."
            baseline_model = None
        except Exception as e:
            logger.error(f"Failed during Ternary RandomForest baseline: {e}", exc_info=True)
            report["notes"] = f"Failed with error: {e}"
            baseline_model = None
            
        # Save Report
        if io:
            try:
                io.save_json(report, "baseline_ternary_rf_report", section='results', use_txt=True)
                logger.info("Saved Ternary RandomForest baseline report.")
            except Exception as save_e:
                 logger.error(f"Failed to save Ternary RF report: {save_e}")
        else:
             logger.warning("IOManager not provided, skipping saving Ternary RF report.")

        return baseline_model, report

    def _run_edge_filtered_check(self, 
                                 predictor: Any, 
                                 X_val: pd.DataFrame, 
                                 y_val_ordinal: pd.Series, 
                                 edge_threshold: float, 
                                 baseline_name: str, 
                                 is_ternary: bool,
                                 ci_conf_level: float) -> Dict[str, Any]:
        """
        Helper function to run edge-filtered checks on a trained baseline model or calibrator.
        
        Args:
            predictor: The trained sklearn classifier or CalibratedClassifierCV instance. # Updated docstring
            X_val: The original SCALED validation features.
            y_val_ordinal: The original ORDINAL validation labels (0,1 or 0,1,2).
            edge_threshold: The threshold |p_up - p_down| >= thr.
            baseline_name: Name for logging (e.g., "Binary-LR").
            is_ternary: Flag indicating if the model/labels are ternary.
            ci_conf_level: Confidence level for CI calculation.

        Returns:
            Dictionary with edge-filtered metrics.
        """
        edge_report = {
            "edge_accuracy": np.nan,
            "edge_ci_lower_bound": np.nan,
            "edge_num_samples_val": len(y_val_ordinal),
            "edge_num_samples_val_filtered": 0,
        }

        if predictor is None: # Check predictor
             logger.warning(f"{baseline_name} edge-filtered check skipped: Predictor not available.")
             return edge_report
        if X_val.empty or y_val_ordinal.empty:
            logger.warning(f"{baseline_name} edge-filtered check skipped: Validation data is empty.")
            return edge_report
            
        logger.info(f"Running edge-filtered check for {baseline_name} (Threshold={edge_threshold:.3f})...")

        try:
            # Predict probabilities using the provided predictor (base or calibrated)
            if not hasattr(predictor, "predict_proba"):
                 logger.error(f"{baseline_name} edge-filtered check failed: Predictor has no predict_proba method.")
                 return edge_report

            # --- Handle potential NaNs and Infs in validation data before prediction --- #
            # Assume X_val is clean here, cleaning happens before calling run_checks
            X_val_clean = X_val 
            # --- End Handle NaNs/Infs --- #

            p_pred = predictor.predict_proba(X_val_clean) # Use predictor.predict_proba

            if is_ternary:
                # Ternary edge calculation: P(up) - P(down)
                if p_pred.shape[1] != 3:
                     logger.error(f"{baseline_name} edge check failed: Expected 3 probability columns for ternary, got {p_pred.shape[1]}.")
                     return edge_report
                edge = p_pred[:, 2] - p_pred[:, 0] 
                y_true_binary_for_edge = (y_val_ordinal == 2).astype(int) 
            else: # Binary case
                # Binary edge calculation: P(up) - P(down) = 2 * P(up) - 1
                if p_pred.shape[1] != 2:
                     logger.error(f"{baseline_name} edge check failed: Expected 2 probability columns for binary, got {p_pred.shape[1]}.")
                     return edge_report
                p_up = p_pred[:, 1]
                edge = 2 * p_up - 1 
                y_true_binary_for_edge = y_val_ordinal.astype(int) # Already 0 or 1

            # Filter samples based on edge threshold
            mask_edge = np.abs(edge) >= edge_threshold
            n_filtered = int(mask_edge.sum())
            edge_report["edge_num_samples_val_filtered"] = n_filtered

            if n_filtered > 0:
                # Evaluate accuracy on the filtered subset
                y_true_filtered = y_true_binary_for_edge[mask_edge]
                edge_pred_dir = (edge[mask_edge] > 0).astype(int) # Predict 1 if edge > 0, else 0
                
                edge_acc = accuracy_score(y_true_filtered, edge_pred_dir)
                edge_report["edge_accuracy"] = float(edge_acc)

                # Calculate Edge CI lower bound (always vs 0.5 chance for binary direction)
                k_correct_edge = np.sum(y_true_filtered == edge_pred_dir)
                edge_ci_lb = self._calculate_ci_lower_bound(k_correct_edge, n_filtered, p=0.5, conf_level=ci_conf_level)
                edge_report["edge_ci_lower_bound"] = edge_ci_lb
                logger.info(f"{baseline_name} Edge-Filtered: Acc={edge_acc:.3f}, CI_LB={edge_ci_lb:.3f}, N_filt={n_filtered}/{edge_report['edge_num_samples_val']}")
            else:
                logger.warning(f"{baseline_name} Edge-Filtered: No validation samples met edge threshold ({edge_threshold}). Edge metrics are NaN.")

        except Exception as edge_err:
            logger.error(f"Error during {baseline_name} edge-filtered calculation: {edge_err}", exc_info=True)

        return edge_report

    # <<< NEW METHOD: Forward Baseline Check >>>
    def run_forward_baseline_check(
        self,
        X_train_fold: pd.DataFrame,
        y_train_fold_ordinal: pd.Series,
        X_test_fold: pd.DataFrame,
        y_test_fold_ordinal: Optional[pd.Series],
        fold_num: int,
        io: Optional[IOManager] = None
    ) -> bool:
        """
        Runs a forward-looking baseline check using the fold's train/test split.

        Trains a Binary Logistic Regression on the entire training fold's data
        and evaluates its CI lower bound on the test fold's data.

        Args:
            X_train_fold: SCALED training features for the fold.
            y_train_fold_ordinal: ORDINAL {0, 1, 2} training labels for the fold.
            X_test_fold: SCALED test features for the fold.
            y_test_fold_ordinal: Optional ORDINAL {0, 1, 2} test labels for the fold.
            fold_num: Current fold number.
            io: IOManager instance (optional, for saving detailed report).

        Returns:
            bool: True if the forward baseline check passes the threshold, False otherwise.
        """
        logger.info(f"--- Fold {fold_num}: Running Forward Baseline Check --- ")
        gate_cfg = self.config.get('validation_gates', {})
        fwd_gate_cfg = gate_cfg.get('forward_baseline', {}) # New config section
        run_fwd_check = fwd_gate_cfg.get('enabled', True) # Check if enabled
        ci_threshold = fwd_gate_cfg.get('ci_threshold', 0.52) # Threshold for this check
        
        if not run_fwd_check:
             logger.info(f"Fold {fold_num}: Skipping forward baseline check as per config.")
             return True # Pass if disabled

        # --- Check Prerequisites ---
        if X_train_fold is None or X_train_fold.empty or \
           y_train_fold_ordinal is None or y_train_fold_ordinal.empty or \
           X_test_fold is None or X_test_fold.empty or \
           y_test_fold_ordinal is None or y_test_fold_ordinal.empty:
            logger.error(f"Fold {fold_num}: Missing or empty train/test data/labels for forward baseline check. Check Failed.")
            return False

        # --- Prepare Binary Labels --- #
        try:
            y_train_fold_binary = y_train_fold_ordinal.map({0: 0, 1: 1, 2: 1})
            y_test_fold_binary = y_test_fold_ordinal.map({0: 0, 1: 1, 2: 1})
            if y_train_fold_binary.isnull().any() or y_test_fold_binary.isnull().any():
                raise ValueError("NaNs found after mapping ordinal to binary labels.")
        except Exception as e:
             logger.error(f"Fold {fold_num}: Error deriving binary labels for forward check: {e}. Check Failed.")
             return False

        # --- Clean Data (NaN/Inf) --- #
        try:
            X_train_clean = pd.DataFrame(np.nan_to_num(X_train_fold, nan=0.0, posinf=0.0, neginf=0.0), index=X_train_fold.index, columns=X_train_fold.columns)
            X_test_clean = pd.DataFrame(np.nan_to_num(X_test_fold, nan=0.0, posinf=0.0, neginf=0.0), index=X_test_fold.index, columns=X_test_fold.columns)
            if X_train_clean.isnull().values.any() or np.isinf(X_train_clean).values.any() or \
               X_test_clean.isnull().values.any() or np.isinf(X_test_clean).values.any():
                 logger.warning(f"Fold {fold_num}: NaNs/Infs persisted after cleaning in forward check.")
        except Exception as e:
            logger.error(f"Fold {fold_num}: Error cleaning data for forward check: {e}. Check Failed.")
            return False
            
        # --- Train Model --- #
        model = None
        try:
            baseline_cfg = self.config.get('baselines', {}).get('logistic_regression', {})
            max_iter = baseline_cfg.get('max_iter', 1000)
            solver = baseline_cfg.get('solver', 'lbfgs')
            random_state = baseline_cfg.get('random_state', 42)
            
            logger.info(f"Fold {fold_num}: Fitting Binary LogReg on full training fold ({len(y_train_fold_binary)} samples) for forward check...")
            model = LogisticRegression(
                max_iter=max_iter,
                solver=solver,
                random_state=random_state
            )
            model.fit(X_train_clean, y_train_fold_binary)
        except Exception as e:
            logger.error(f"Fold {fold_num}: Failed to train LogReg model for forward check: {e}", exc_info=True)
            return False # Fail check if model training fails
            
        # --- Evaluate on Test Set & Calculate CI LB --- #
        ci_lb = np.nan
        accuracy = np.nan
        n_test = len(y_test_fold_binary)
        status = "FAIL"
        try:
            logger.debug(f"Fold {fold_num}: Predicting on test fold ({n_test} samples) for forward check...")
            y_pred_test = model.predict(X_test_clean)
            accuracy = accuracy_score(y_test_fold_binary, y_pred_test)
            k_correct = np.sum(y_pred_test == y_test_fold_binary)
            
            # Get confidence level from config (reuse from binary baseline section?)
            bin_baseline_cfg = self.config.get('validation_gates', {}).get('baseline_binary', {})
            ci_conf_level = bin_baseline_cfg.get('ci_confidence_level', 0.95)
            
            ci_lb = self._calculate_ci_lower_bound(k_correct, n_test, p=0.5, conf_level=ci_conf_level)
            
            if pd.isna(ci_lb):
                 logger.warning(f"Forward Baseline Check (Fold {fold_num}): CI LB calculation resulted in NaN (Acc={accuracy:.4f}, N={n_test}). Check Failed.")
                 status = "FAIL (NaN)"
            elif ci_lb >= ci_threshold:
                 logger.info(f"Forward Baseline Check PASSED (Fold {fold_num}): Test CI LB {ci_lb:.4f} >= {ci_threshold} (Acc={accuracy:.4f}, N={n_test})")
                 status = "PASS"
            else:
                 logger.error(f"Forward Baseline Check FAILED (Fold {fold_num}): Test CI LB {ci_lb:.4f} < {ci_threshold} (Acc={accuracy:.4f}, N={n_test})")
                 status = "FAIL"
                 
        except Exception as e:
            logger.error(f"Fold {fold_num}: Error during prediction or CI calculation for forward check: {e}", exc_info=True)
            status = "FAIL (Error)"

        # --- Save Detailed Report (Optional) --- #
        if io:
            report = {
                "fold_num": fold_num,
                "check_type": "forward_baseline_logistic_regression",
                "n_train": len(y_train_fold_binary),
                "n_test": n_test,
                "test_accuracy": accuracy,
                "test_ci_lower_bound": ci_lb,
                "ci_threshold": ci_threshold,
                "status": status
            }
            try:
                io.save_json(report, f"forward_baseline_report_fold_{fold_num}", section='results', use_txt=True)
            except Exception as save_e:
                logger.error(f"Failed to save forward baseline report for fold {fold_num}: {save_e}")

        return status == "PASS"
    # <<< END NEW METHOD >>>

    # --- Main Check Runner --- #
    def run_checks(self, 
                   X_train: pd.DataFrame, 
                   y_train_ordinal: pd.Series, 
                   X_val: pd.DataFrame, 
                   y_val_ordinal: Optional[pd.Series], # Now always ordinal (0,1 or 0,1,2)
                   fold_num: int, 
                   io_manager: Optional[IOManager] 
                   ) -> bool: # Returns True if checks allow pipeline continuation
        """ 
        Runs ALL enabled baseline checks (binary and ternary, respecting individual flags).

        Args:
            X_train: SCALED training features.
            y_train_ordinal: ORDINAL training labels {0, 1} or {0, 1, 2}.
            X_val: SCALED validation features.
            y_val_ordinal: ORDINAL validation labels {0, 1} or {0, 1, 2}. CAN BE NONE if val set is empty.
            fold_num: Current fold number for logging/saving.
            io_manager: IOManager instance.

        Returns:
            bool: True if all enabled and required checks passed thresholds, False otherwise.
                  Returns True if master baseline check flag is disabled.
        """
        logger.info(f"--- Fold {fold_num}: Running Baseline Checks (Binary & Ternary if enabled) ---")
        overall_checks_passed = True # Assume passing unless a gate fails
        check_reports = [] # Store reports from individual baselines
        summary_data = [] # <<< ADDED: List to store data for the summary table

        # --- Check if Validation Labels Exist --- #
        val_labels_exist = y_val_ordinal is not None and not y_val_ordinal.empty
        if not val_labels_exist:
             logger.warning(f"Fold {fold_num}: Validation labels (y_val_ordinal) are missing or empty. Checks requiring validation data will be skipped or may fail.")

        # --- Determine if checks should run at all (Master Switch) --- #
        gate_cfg = self.config.get('validation_gates', {})
        run_baseline_check = gate_cfg.get('run_baseline_check', True)
        if not run_baseline_check:
             logger.info(f"Fold {fold_num}: Skipping baseline checks as per master config (validation_gates.run_baseline_check = False).")
             return True

        # --- Prepare Binary Labels (Derive from Ordinal) --- #
        # Mapping: 0 -> 0 (Down), 1 -> 1 (Flat -> Up), 2 -> 1 (Up -> Up)
        y_train_binary = y_train_ordinal.map({0: 0, 1: 1, 2: 1}) if y_train_ordinal is not None else None
        y_val_binary = y_val_ordinal.map({0: 0, 1: 1, 2: 1}) if val_labels_exist else None
        if y_train_binary is None:
            logger.error(f"Fold {fold_num}: Failed to derive binary training labels. Aborting baseline checks.")
            raise SystemExit(f"Fold {fold_num}: Failed to derive binary training labels for baseline checks.")
        logger.debug(f"Fold {fold_num}: Derived binary labels for checks.")
        # --- End Prepare Binary Labels --- #

        # --- Initialize flags for the new "either/or" raw gate logic --- #
        lr_raw_passed = False
        rf_raw_passed = False
        # --- Initialize flags for the new "either/or" edge gate logic --- #
        lr_edge_passed = False
        rf_edge_passed = False

        # --- Run BINARY CHECKS (Baselines 1 & 2) --- #
        logger.info(f"--- Fold {fold_num}: Evaluating BINARY Baselines ---")
        binary_gate_cfg = gate_cfg.get('baseline_binary', {})
        binary_ci_threshold = binary_gate_cfg.get('ci_threshold', 0.51)
        # --- Add config for edge-filtered check --- #
        binary_edge_threshold_value = binary_gate_cfg.get('edge_threshold_value', 0.1) # Threshold for *calculating* edge metrics
        binary_edge_ci_threshold_gate = binary_gate_cfg.get('edge_ci_threshold_gate', 0.60) # Threshold for the *gate*
        run_binary_lr_edge_check = binary_gate_cfg.get('run_logistic_regression_edge_check', True) # Specific flag for edge check
        # --- End Add config --- #
        run_binary_lr = binary_gate_cfg.get('run_logistic_regression', True)
        run_binary_rf = binary_gate_cfg.get('run_random_forest', False)

        # Baseline 1: Binary Logistic Regression
        if run_binary_lr:
            logger.info(f"--- Fold {fold_num}: Running Binary Logistic Regression (Baseline 1) ---")
            y_val_binary_input = y_val_binary if val_labels_exist else pd.Series(dtype=int)
            # <<< Capture both base_lr and final_predictor >>>
            lr_model, lr_predictor, lr_report = self.run_binary_logistic_baseline(X_train, y_train_binary, X_val, y_val_binary_input, self.io)
            check_reports.append(lr_report)
            
            # --- Raw CI Check (Failure Gate - uses BASE model results stored in report) --- #
            raw_ci_status = "SKIPPED (NaN)"
            raw_ci_value = np.nan
            if 'ci_lower_bound' in lr_report and not pd.isna(lr_report['ci_lower_bound']):
                raw_ci_value = lr_report['ci_lower_bound']
                if raw_ci_value < binary_ci_threshold:
                    # overall_checks_passed = False # <<< MODIFIED: Don't fail overall gate here
                    raw_ci_status = "FAIL (Threshold)"
                else:
                    lr_raw_passed = True # <<< ADDED: Set flag if passed
                    raw_ci_status = "PASS (Threshold)"
            else:
                logger.warning(f"BASELINE CHECK SKIPPED (Fold {fold_num}): Binary LogReg Raw CI LB not calculated or NaN. Cannot determine pass/fail.")
                # overall_checks_passed = False # Treat missing mandatory check as failure -> This will be handled by the new gate
                raw_ci_status = "SKIPPED (NaN)"
            
            summary_data.append({
                "Baseline": "Binary LogReg (Raw)", "Metric": "CI Lower Bound", 
                "Value": f"{raw_ci_value:.4f}", "Threshold": f">= {binary_ci_threshold:.2f}", 
                "Status": raw_ci_status, 
                "Is Gate": "See Combined" # <<< MODIFIED: Indicate combined gate
            })
            
            # --- Edge-Filtered Check (Failure Gate - uses FINAL predictor) --- #
            edge_ci_status = "SKIPPED (Config)"
            edge_ci_value = np.nan
            if run_binary_lr_edge_check:
                if overall_checks_passed: # Only run if raw check passed
                    # Use lr_predictor (which is either base LR or CalibratedClassifierCV)
                    if lr_predictor is not None and not X_val.empty and val_labels_exist: 
                        logger.info(f"--- Fold {fold_num}: Running Binary LogReg Edge-Filtered Check (Gate Threshold={binary_edge_ci_threshold_gate:.2f}) ---")
                        # Get confidence level from config (used inside helper)
                        baseline_lr_cfg = self.config.get('baselines', {}).get('logistic_regression', {})
                        ci_conf_level = baseline_lr_cfg.get('ci_confidence_level', 0.95)

                        # <<< Pass lr_predictor to edge check >>>
                        edge_report_lr = self._run_edge_filtered_check(
                            predictor=lr_predictor, # Pass the final predictor (calibrated or not)
                            X_val=X_val, # Pass original scaled X_val (assumed clean)
                            y_val_ordinal=y_val_binary_input, # Pass binary labels (0/1)
                            edge_threshold=binary_edge_threshold_value, # Use config value for calculation
                            baseline_name="Binary LogReg",
                            is_ternary=False,
                            ci_conf_level=ci_conf_level
                        )
                        check_reports.append(edge_report_lr)

                        if 'edge_ci_lower_bound' in edge_report_lr and not pd.isna(edge_report_lr['edge_ci_lower_bound']):
                            edge_ci_value = edge_report_lr['edge_ci_lower_bound']
                            if edge_ci_value < binary_edge_ci_threshold_gate:
                                # overall_checks_passed = False # <<< MODIFIED: Don't fail overall gate here
                                edge_ci_status = "FAIL (Threshold)"
                            else: # Correctly indented else for `if edge_ci_value < ...`
                                lr_edge_passed = True # <<< ADDED: Set flag if passed
                                edge_ci_status = "PASS (Threshold)"
                        else: # Correctly indented else for `if 'edge_ci_lower_bound' in ...`
                            logger.warning(f"BASELINE CHECK SKIPPED (Fold {fold_num}): Binary LogReg Edge CI LB not calculated or NaN. Cannot determine pass/fail.")
                            # overall_checks_passed = False # Treat missing mandatory check as failure -> Handled by combined gate
                            edge_ci_status = "SKIPPED (NaN)" # Added status update
                    else: # Correctly indented else for `if lr_predictor is not None ...`
                        logger.warning(f"Skipping Binary LogReg Edge-Filtered check: Predictor not available, validation data empty, or validation labels missing.")
                        # overall_checks_passed = False # Consider it a failure if mandatory edge check couldn't run -> Handled by combined gate
                        edge_ci_status = "SKIPPED (Prereq)" # Added status update
                else: # Correctly indented else for `if overall_checks_passed:`
                    logger.info(f"Skipping Binary LogReg Edge-Filtered check as Raw CI check already failed.")
                    edge_ci_status = "SKIPPED (Prior Fail)"
            # else: # run_binary_lr_edge_check is False, status remains SKIPPED (Config)

            summary_data.append({
                "Baseline": "Binary LogReg (Edge)", "Metric": "CI Lower Bound",
                "Value": f"{edge_ci_value:.4f}", "Threshold": f">= {binary_edge_ci_threshold_gate:.2f}",
                "Status": edge_ci_status, 
                "Is Gate": "See Combined Edge" if run_binary_lr_edge_check else "No" # <<< MODIFIED
            })
            # --- End Edge-Filtered Check --- #
        else: # Correctly indented else for `if run_binary_lr:`
            logger.info(f"Fold {fold_num}: Skipping Binary Logistic Regression baseline as per config.")
            summary_data.append({"Baseline": "Binary LogReg (Raw)", "Metric": "CI Lower Bound", "Value": "N/A", "Threshold": "N/A", "Status": "SKIPPED (Config)", "Is Gate": True})
            summary_data.append({"Baseline": "Binary LogReg (Edge)", "Metric": "CI Lower Bound", "Value": "N/A", "Threshold": "N/A", "Status": "SKIPPED (Config)", "Is Gate": True}) # Gate status depends on run_binary_lr_edge_check, but if LR is skipped, edge is too

        # Baseline 2: Binary Random Forest (Monitor Only -> ABORT if enabled)
        rf_model = None # Initialize rf_model
        rf_ci_status = "SKIPPED (Config)"
        rf_ci_value = np.nan
        rf_edge_ci_status = "SKIPPED (Config)"
        rf_edge_ci_value = np.nan

        # --- Corrected Config Key: Use same section 'baseline_binary' --- #
        run_binary_rf = binary_gate_cfg.get('run_random_forest', False)
        binary_rf_ci_threshold = binary_gate_cfg.get('binary_rf_ci_lb', 0.54) # G2 Threshold
        run_binary_rf_edge_check = binary_gate_cfg.get('run_random_forest_edge_check', False) # G9 Flag
        binary_rf_edge_ci_threshold_gate = binary_gate_cfg.get('edge_binary_rf_ci_lb', 0.62) # G9 Threshold
        # --- End Config Key Correction --- #

        if run_binary_rf:
            logger.info(f"--- Fold {fold_num}: Running Binary Random Forest (Baseline 2) ---")
            y_val_binary_input = y_val_binary if val_labels_exist else pd.Series(dtype=int)
            rf_model, rf_report = self.run_binary_random_forest_baseline(X_train, y_train_binary, X_val, y_val_binary_input, self.io)
            check_reports.append(rf_report)
            if 'ci_lower_bound' in rf_report and not pd.isna(rf_report['ci_lower_bound']):
                rf_ci_value = rf_report['ci_lower_bound']
                # --- G2: Abort if enabled and fails threshold --- #
                if rf_ci_value < binary_rf_ci_threshold:
                    # logger.error(...) # Keep as error log but also trigger failure
                    # logger.error(f"BASELINE FAILED (GATE): Binary RF Raw CI LB {rf_ci_value:.4f} < {binary_rf_ci_threshold}")
                    # overall_checks_passed = False # <<< MODIFIED: Don't fail overall gate here (G2)
                    rf_ci_status = "FAIL (Threshold)"
                else: # Correctly indented else for `if rf_ci_value < ...`
                    rf_raw_passed = True # <<< ADDED: Set flag if passed
                    rf_ci_status = "PASS (Threshold)"
            else: # Correctly indented else for `if 'ci_lower_bound' in ...`
                logger.warning(f"BASELINE CHECK SKIPPED (Fold {fold_num}): Binary RF Raw CI LB not calculated or NaN. Cannot determine pass/fail.")
                # overall_checks_passed = False # <<< MODIFIED: Don't fail overall gate here (G2)
                rf_ci_status = "SKIPPED (NaN)"

            # --- G9: Edge Binary RF Check (Abort if enabled) --- #
            if run_binary_rf_edge_check:
                if overall_checks_passed: # Only run if raw check passed
                    if rf_model is not None and not X_val.empty and val_labels_exist:
                        logger.info(f"--- Fold {fold_num}: Running Binary RF Edge-Filtered Check (Gate Threshold={binary_rf_edge_ci_threshold_gate:.2f}) ---")
                        baseline_rf_cfg = self.config.get('baselines', {}).get('random_forest', {})
                        ci_conf_level = baseline_rf_cfg.get('ci_confidence_level', 0.95)

                        edge_report_rf = self._run_edge_filtered_check(
                            predictor=rf_model, # <<< Fix: Changed model to predictor
                            X_val=X_val,
                            y_val_ordinal=y_val_binary_input,
                            edge_threshold=binary_edge_threshold_value,
                            baseline_name="Binary RF",
                            is_ternary=False,
                            ci_conf_level=ci_conf_level
                        )
                        check_reports.append(edge_report_rf)

                        if 'edge_ci_lower_bound' in edge_report_rf and not pd.isna(edge_report_rf['edge_ci_lower_bound']):
                            rf_edge_ci_value = edge_report_rf['edge_ci_lower_bound']
                            if rf_edge_ci_value < binary_rf_edge_ci_threshold_gate:
                                # overall_checks_passed = False # <<< MODIFIED: Don't fail overall gate here (G9)
                                rf_edge_ci_status = "FAIL (Threshold)"
                            else: # Correctly indented else for `if rf_edge_ci_value < ...`
                                rf_edge_passed = True # <<< ADDED: Set flag if passed
                                rf_edge_ci_status = "PASS (Threshold)"
                        else: # Correctly indented else for `if 'edge_ci_lower_bound' in ...`
                            logger.warning(f"BASELINE CHECK SKIPPED (Fold {fold_num}): Binary RF Edge CI LB not calculated or NaN. Cannot determine pass/fail.") # Corrected log message
                            # overall_checks_passed = False # Fail if mandatory G9 cannot run -> Handled by combined gate
                            rf_edge_ci_status = "SKIPPED (NaN)"
                    else: # Correctly indented else for `if rf_model is not None ...`
                        logger.warning(f"Skipping Binary RF Edge-Filtered check: Model not trained, validation data empty, or validation labels missing.")
                        # overall_checks_passed = False # Fail if mandatory G9 cannot run -> Handled by combined gate
                        rf_edge_ci_status = "SKIPPED (Prereq)" # Added status update
                else: # Correctly indented else for `if overall_checks_passed:`
                    logger.info(f"Skipping Binary RF Edge-Filtered check as Raw CI check already failed.")
                    rf_edge_ci_status = "SKIPPED (Prior Fail)"
            # else: # run_binary_rf_edge_check is False, status remains SKIPPED (Config)
            # --- End G9 --- #
        else: # Correctly indented else for `if run_binary_rf:`
            # If RF baseline is not run, status remains SKIPPED (Config) for both raw and edge
            pass # No action needed, statuses already initialized

        summary_data.append({ # Add RF Raw to summary
            "Baseline": "Binary RF (Raw)", "Metric": "CI Lower Bound",
            "Value": f"{rf_ci_value:.4f}", "Threshold": f">= {binary_rf_ci_threshold:.2f}",
            "Status": rf_ci_status, 
            "Is Gate": "See Combined" # <<< MODIFIED: Indicate combined gate
        })
        summary_data.append({ # Add RF Edge to summary
            "Baseline": "Binary RF (Edge)", "Metric": "CI Lower Bound",
            "Value": f"{rf_edge_ci_value:.4f}", "Threshold": f">= {binary_rf_edge_ci_threshold_gate:.2f}",
            "Status": rf_edge_ci_status, "Is Gate": "See Combined Edge" if run_binary_rf_edge_check else "No" # <<< MODIFIED
        })
        # --- End BINARY CHECKS --- #

        # --- New Gate: Abort if NEITHER Raw LR NOR Raw RF passed --- #
        # This logic applies if *any* binary raw check was configured to run.
        # If both were disabled, this gate is bypassed.
        lr_check_enabled = binary_gate_cfg.get('run_logistic_regression', True)
        rf_check_enabled = binary_gate_cfg.get('run_random_forest', False)
        if (lr_check_enabled or rf_check_enabled) and not (lr_raw_passed or rf_raw_passed):
            logger.error(f"BASELINE FAILED (COMBINED GATE): Neither Binary LogReg Raw ({'PASS' if lr_raw_passed else 'FAIL/SKIPPED'}) nor Binary RF Raw ({'PASS' if rf_raw_passed else 'FAIL/SKIPPED'}) passed their respective thresholds. At least one must pass. Aborting fold.")
            overall_checks_passed = False # Set overall flag to ensure the fold aborts

        # --- New Gate: Abort if NEITHER Enabled Edge LR NOR Enabled Edge RF passed --- #
        lr_edge_check_enabled = binary_gate_cfg.get('run_logistic_regression_edge_check', True)
        rf_edge_check_enabled = binary_gate_cfg.get('run_random_forest_edge_check', False)
        if (lr_edge_check_enabled or rf_edge_check_enabled) and not (lr_edge_passed or rf_edge_passed):
             logger.error(f"BASELINE FAILED (COMBINED EDGE GATE): Neither Binary LogReg Edge ({'PASS' if lr_edge_passed else 'FAIL/SKIPPED'}) nor Binary RF Edge ({'PASS' if rf_edge_passed else 'FAIL/SKIPPED'}) passed their respective thresholds (when enabled). At least one enabled edge check must pass. Aborting fold.")
             overall_checks_passed = False # Set overall flag to ensure the fold aborts

        # --- Run TERNARY CHECKS (Baselines 3 & 4) - Monitor Only (unless edge check fails) --- #
        # Initialize ternary statuses
        ternary_lr_ci_status = "SKIPPED (Config)"
        ternary_lr_ci_value = np.nan
        ternary_rf_ci_status = "SKIPPED (Config)"
        ternary_rf_ci_value = np.nan
        ternary_lr_edge_ci_status = "SKIPPED (Config)"
        ternary_lr_edge_ci_value = np.nan
        ternary_rf_edge_ci_status = "SKIPPED (Config)"
        ternary_rf_edge_ci_value = np.nan
        ternary_labels_missing = False
        lr_model_t = None # Initialize ternary models
        rf_model_t = None

        use_ternary_config = self.config.get('gru', {}).get('use_ternary', False)
        if use_ternary_config:
            logger.info(f"--- Fold {fold_num}: Evaluating TERNARY Baselines (use_ternary=True in config) --- ")
            ternary_gate_cfg = gate_cfg.get('baseline_ternary', {})
            # --- Add config keys for Ternary checks --- #
            ternary_ci_threshold = ternary_gate_cfg.get('ternary_ci_lb', 0.40) # G3 Threshold
            ternary_rf_ci_threshold = ternary_gate_cfg.get('ternary_rf_ci_lb', 0.42) # G4 Threshold
            run_ternary_lr = ternary_gate_cfg.get('run_logistic_regression', True)
            run_ternary_rf = ternary_gate_cfg.get('run_random_forest', False)
            run_ternary_lr_edge_check = ternary_gate_cfg.get('run_logistic_regression_edge_check', False) # G10 Flag
            ternary_edge_threshold_value = ternary_gate_cfg.get('edge_threshold_value', 0.1)
            ternary_lr_edge_ci_threshold_gate = ternary_gate_cfg.get('edge_ternary_ci_lb', 0.57) # G10 Threshold
            run_ternary_rf_edge_check = ternary_gate_cfg.get('run_random_forest_edge_check', False) # G11 Flag
            ternary_rf_edge_ci_threshold_gate = ternary_gate_cfg.get('edge_ternary_rf_ci_lb', 0.58) # G11 Threshold
            # --- End Add config keys --- #

            if not val_labels_exist:
                logger.error(f"Fold {fold_num}: Cannot run Ternary checks because validation labels are missing.")
                ternary_labels_missing = True
                # NOTE: overall_checks_passed = False removed - handled by final check if mandatory
                ternary_lr_ci_status = "SKIPPED (Labels Missing)"
                ternary_rf_ci_status = "SKIPPED (Labels Missing)"
                ternary_lr_edge_ci_status = "SKIPPED (Labels Missing)"
                ternary_rf_edge_ci_status = "SKIPPED (Labels Missing)"
            else:
                # Baseline 3: Ternary Logistic Regression (Raw: Warn, Edge: Warn)
                if run_ternary_lr:
                    logger.info(f"--- Fold {fold_num}: Running Ternary Logistic Regression (Baseline 3) ---")
                    lr_model_t, lr_report_t = self.run_ternary_logistic_baseline(X_train, y_train_ordinal, X_val, y_val_ordinal, self.io)
                    check_reports.append(lr_report_t)
                    if 'ci_lower_bound' in lr_report_t and not pd.isna(lr_report_t['ci_lower_bound']):
                        ternary_lr_ci_value = lr_report_t['ci_lower_bound']
                        if ternary_lr_ci_value < ternary_ci_threshold:
                            logger.warning(f"MONITOR (BASELINE WARN): Ternary LogReg Raw CI LB {ternary_lr_ci_value:.4f} < {ternary_ci_threshold}")
                            ternary_lr_ci_status = "WARN (Low Raw CI)"
                        else: # Correctly indented else for `if ternary_lr_ci_value < ...`
                            ternary_lr_ci_status = "PASS (Raw CI)"
                    else: # Correctly indented else for `if 'ci_lower_bound' in ...`
                        logger.warning(f"MONITOR (BASELINE SKIPPED): Ternary LogReg Raw CI LB not calculated or NaN.")
                        ternary_lr_ci_status = "SKIPPED (NaN)"

                    # --- G10: Edge Ternary LR Check (Warn) --- #
                    if run_ternary_lr_edge_check: # Correctly indented if
                        if lr_model_t is not None:
                            logger.info(f"--- Fold {fold_num}: Running Ternary LogReg Edge-Filtered Check (Warn Threshold={ternary_lr_edge_ci_threshold_gate:.2f}) ---")
                            baseline_tlr_cfg = self.config.get('baselines', {}).get('multinomial_logistic_regression', {})
                            ci_conf_level = baseline_tlr_cfg.get('ci_confidence_level', 0.95)

                            edge_report_tlr = self._run_edge_filtered_check(
                                predictor=lr_model_t, # <<< Fix: Changed model to predictor
                                X_val=X_val,
                                y_val_ordinal=y_val_ordinal, # Pass original ternary labels
                                edge_threshold=ternary_edge_threshold_value,
                                baseline_name="Ternary LogReg",
                                is_ternary=True,
                                ci_conf_level=ci_conf_level
                            )
                            check_reports.append(edge_report_tlr)
                            if 'edge_ci_lower_bound' in edge_report_tlr and not pd.isna(edge_report_tlr['edge_ci_lower_bound']):
                                ternary_lr_edge_ci_value = edge_report_tlr['edge_ci_lower_bound']
                                if ternary_lr_edge_ci_value < ternary_lr_edge_ci_threshold_gate:
                                    logger.warning(f"MONITOR (BASELINE WARN): Ternary LogReg Edge CI LB {ternary_lr_edge_ci_value:.4f} < {ternary_lr_edge_ci_threshold_gate}")
                                    ternary_lr_edge_ci_status = "WARN (Low Edge CI)"
                                else: # Correctly indented else for `if ternary_lr_edge_ci_value < ...`
                                    ternary_lr_edge_ci_status = "PASS (Edge CI)"
                            else: # Correctly indented else for `if 'edge_ci_lower_bound' in ...`
                                logger.warning(f"MONITOR (BASELINE SKIPPED): Ternary LogReg Edge CI LB not calculated or NaN.")
                                ternary_lr_edge_ci_status = "SKIPPED (NaN)"
                        else: # Correctly indented else for `if lr_model_t is not None:`
                            logger.warning(f"Skipping Ternary LogReg Edge-Filtered check: Model not trained.")
                            ternary_lr_edge_ci_status = "SKIPPED (Model)"
                    else: # Correctly indented else for `if run_ternary_lr_edge_check:`
                        ternary_lr_edge_ci_status = "SKIPPED (Config)"
                    # --- End G10 --- #
                else: # Correctly indented else for `if run_ternary_lr:`
                    ternary_lr_ci_status = "SKIPPED (Config)"
                    ternary_lr_edge_ci_status = "SKIPPED (Config)"

                # Baseline 4: Ternary Random Forest (Raw: Warn, Edge: Warn)
                if run_ternary_rf:
                    logger.info(f"--- Fold {fold_num}: Running Ternary Random Forest (Baseline 4) ---")
                    rf_model_t, rf_report_t = self.run_ternary_random_forest_baseline(X_train, y_train_ordinal, X_val, y_val_ordinal, self.io)
                    check_reports.append(rf_report_t)
                    if 'ci_lower_bound' in rf_report_t and not pd.isna(rf_report_t['ci_lower_bound']):
                        ternary_rf_ci_value = rf_report_t['ci_lower_bound']
                        if ternary_rf_ci_value < ternary_rf_ci_threshold: # Use G4 threshold
                            logger.warning(f"MONITOR (BASELINE WARN): Ternary RF Raw CI LB {ternary_rf_ci_value:.4f} < {ternary_rf_ci_threshold}")
                            ternary_rf_ci_status = "WARN (Low Raw CI)"
                        else: # Correctly indented else for `if ternary_rf_ci_value < ...`
                            ternary_rf_ci_status = "PASS (Raw CI)"
                    else: # Correctly indented else for `if 'ci_lower_bound' in ...`
                        logger.warning(f"MONITOR (BASELINE SKIPPED): Ternary RF Raw CI LB not calculated or NaN.")
                        ternary_rf_ci_status = "SKIPPED (NaN)"

                    # --- G11: Edge Ternary RF Check (Warn) --- #
                    if run_ternary_rf_edge_check: # Correctly indented if
                        if rf_model_t is not None:
                            logger.info(f"--- Fold {fold_num}: Running Ternary RF Edge-Filtered Check (Warn Threshold={ternary_rf_edge_ci_threshold_gate:.2f}) ---")
                            baseline_trf_cfg = self.config.get('baselines', {}).get('ternary_random_forest', {})
                            ci_conf_level = baseline_trf_cfg.get('ci_confidence_level', 0.95)

                            edge_report_trf = self._run_edge_filtered_check(
                                predictor=rf_model_t, # <<< Fix: Changed model to predictor
                                X_val=X_val,
                                y_val_ordinal=y_val_ordinal,
                                edge_threshold=ternary_edge_threshold_value,
                                baseline_name="Ternary RF",
                                is_ternary=True,
                                ci_conf_level=ci_conf_level
                            )
                            check_reports.append(edge_report_trf)
                            if 'edge_ci_lower_bound' in edge_report_trf and not pd.isna(edge_report_trf['edge_ci_lower_bound']):
                                ternary_rf_edge_ci_value = edge_report_trf['edge_ci_lower_bound']
                                if ternary_rf_edge_ci_value < ternary_rf_edge_ci_threshold_gate:
                                    logger.warning(f"MONITOR (BASELINE WARN): Ternary RF Edge CI LB {ternary_rf_edge_ci_value:.4f} < {ternary_rf_edge_ci_threshold_gate}")
                                    ternary_rf_edge_ci_status = "WARN (Low Edge CI)"
                                else: # Correctly indented else for `if ternary_rf_edge_ci_value < ...`
                                    ternary_rf_edge_ci_status = "PASS (Edge CI)"
                            else: # Correctly indented else for `if 'edge_ci_lower_bound' in ...`
                                logger.warning(f"MONITOR (BASELINE SKIPPED): Ternary RF Edge CI LB not calculated or NaN.")
                                ternary_rf_edge_ci_status = "SKIPPED (NaN)"
                        else: # Correctly indented else for `if rf_model_t is not None:`
                            logger.warning(f"Skipping Ternary RF Edge-Filtered check: Model not trained.")
                            ternary_rf_edge_ci_status = "SKIPPED (Model)"
                    else: # Correctly indented else for `if run_ternary_rf_edge_check:`
                        ternary_rf_edge_ci_status = "SKIPPED (Config)"
                    # --- End G11 --- #
                else: # Correctly indented else for `if run_ternary_rf:`
                    ternary_rf_ci_status = "SKIPPED (Config)"
                    ternary_rf_edge_ci_status = "SKIPPED (Config)"

        # Add ternary results to summary table
        summary_data.append({ # Ternary LR Raw
            "Baseline": "Ternary LogReg (Raw)", "Metric": "CI Lower Bound",
            "Value": f"{ternary_lr_ci_value:.4f}", "Threshold": f">= {ternary_ci_threshold:.2f}",
            "Status": ternary_lr_ci_status, "Is Gate": False # G3: Warn
        })
        summary_data.append({ # Ternary LR Edge
            "Baseline": "Ternary LogReg (Edge)", "Metric": "CI Lower Bound",
            "Value": f"{ternary_lr_edge_ci_value:.4f}", "Threshold": f">= {ternary_lr_edge_ci_threshold_gate:.2f}",
            "Status": ternary_lr_edge_ci_status, "Is Gate": False # G10: Warn
        })
        summary_data.append({ # Ternary RF Raw
            "Baseline": "Ternary RF (Raw)", "Metric": "CI Lower Bound",
            "Value": f"{ternary_rf_ci_value:.4f}", "Threshold": f">= {ternary_rf_ci_threshold:.2f}",
            "Status": ternary_rf_ci_status, "Is Gate": False # G4: Warn
        })
        summary_data.append({ # Ternary RF Edge
            "Baseline": "Ternary RF (Edge)", "Metric": "CI Lower Bound",
            "Value": f"{ternary_rf_edge_ci_value:.4f}", "Threshold": f">= {ternary_rf_edge_ci_threshold_gate:.2f}",
            "Status": ternary_rf_edge_ci_status, "Is Gate": False # G11: Warn
        })
        # --- End TERNARY CHECKS --- #


        # --- Aggregate and Save Overall Report --- #
        overall_report = {
            "fold_num": fold_num,
            "use_ternary_config": use_ternary_config,
            "baseline_checks_passed_mandatory": overall_checks_passed, # Rename key for clarity
            "individual_reports": check_reports,
            "summary_table_data": summary_data # Add summary data to report
        }
        if self.io:
            try:
                self.io.save_json(overall_report, f"baseline_overall_report_fold_{fold_num}", section='results', use_txt=True)
                logger.info(f"Saved baseline overall report for Fold {fold_num}.")
            except Exception as e:
                logger.error(f"Failed to save baseline overall report for Fold {fold_num}: {e}")
        # --- End Aggregation --- #

        # --- Format and Print Summary Table --- #
        # Calculate column widths dynamically (optional, can use fixed widths)
        widths = {
            "Baseline": max(len(row["Baseline"]) for row in summary_data) if summary_data else 15,
            "Metric": max(len(row["Metric"]) for row in summary_data) if summary_data else 15,
            "Value": 10,
            "Threshold": 10,
            "Status": max(len(row["Status"]) for row in summary_data) if summary_data else 18,
            "Is Gate": 8
        }
        total_width = sum(widths.values()) + len(widths) * 3 + 1 # Account for spacing

        # Header
        header = f"| {'Baseline':<{widths['Baseline']}} | {'Metric':<{widths['Metric']}} | {'Value':>{widths['Value']}} | {'Threshold':>{widths['Threshold']}} | {'Status':<{widths['Status']}} | {'Is Gate':<{widths['Is Gate']}} |"
        separator = '-' * total_width

        # Build table string
        table_str = f"\n{separator}\nBASELINE CHECK SUMMARY - Fold {fold_num}\n{separator}\n{header}\n{separator}\n"
        for row in summary_data:
            # Format NaN values consistently
            value_str = "N/A" if pd.isna(row["Value"]) or row["Value"]=='nan' else str(row["Value"])
            threshold_str = str(row["Threshold"])
            status_str = str(row["Status"])
            gate_str = "Yes" if row["Is Gate"] else "No"
            table_str += f"| {row['Baseline']:<{widths['Baseline']}} | {row['Metric']:<{widths['Metric']}} | {value_str:>{widths['Value']}} | {threshold_str:>{widths['Threshold']}} | {status_str:<{widths['Status']}} | {gate_str:<{widths['Is Gate']}} |\n"
        table_str += f"{separator}\n"

        # Print the table via logger
        logger.info(table_str)
        # --- End Format and Print --- #

        # --- Final Check & Raise SystemExit --- #
        # Add ternary label missing condition to mandatory failure check
        mandatory_failed = not overall_checks_passed
        # Check if ternary LR was supposed to run and labels were missing
        ternary_lr_should_run = use_ternary_config and run_ternary_lr
        if ternary_lr_should_run and ternary_labels_missing:
             logger.error(f"Fold {fold_num}: FAILED MANDATORY baseline check due to missing ternary labels when required for Ternary LogReg.")
             mandatory_failed = True

        if mandatory_failed:
            logger.error(f"Fold {fold_num}: FAILED one or more MANDATORY baseline check gates. See summary table and logs.")
            raise SystemExit(f"Fold {fold_num}: FAILED MANDATORY Baseline Check Gate")
        else:
            final_status_msg = f"Fold {fold_num}: PASSED MANDATORY baseline checks."
            # Add note about monitored checks if any failed or warned
            monitored_issues = [r for r in summary_data if not r["Is Gate"] and ("FAIL" in r["Status"] or "WARN" in r["Status"])]
            if monitored_issues:
                 final_status_msg += " (Note: Some monitored baselines failed/warned - see table)."
            else:
                 final_status_msg += " (Monitored baselines OK or skipped)."
            logger.info(final_status_msg)

        return overall_checks_passed