"""
Benchmark script for UACQR method on pre-processed datasets.

This script runs the UACQR (Uncertainty-Aware Conformal Quantile Regression) method
on prepared datasets and compares against CLEAR's method.

Performance Tips:
- Use the --parallel_jobs argument to control parallelism for UACQR's bootstrap models
- For maximum speed, set --parallel_jobs to a value around 4-8 (based on your system)
- Be careful with high values on machines with limited RAM as each bootstrap model requires memory
- For large datasets or many bootstrap samples, use a moderate number of parallel jobs
- To run on one dataset at a time for maximum speed, provide a specific dataset directory
"""

import os
import sys
import glob
import pickle
import argparse
import numpy as np
import pandas as pd
# from sklearn.model_selection import train_test_split # UACQR helper uses this
from tqdm import tqdm
import time
from datetime import datetime
import logging
from utils import setup_logging, StreamToLogger

# Ensure src and UACQR directories are in the path
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(script_dir, '..', '..'))
src_path = os.path.join(project_root, 'src') # Explicitly define src path
if project_root not in sys.path:
    sys.path.append(project_root)
    print(f"Appended project root: {project_root}")
if src_path not in sys.path: # Add src path
    sys.path.append(src_path)
    print(f"Appended src path: {src_path}")

# Add UACQR directory to sys.path
uacqr_base_path = os.path.join(project_root, "UACQR")
if os.path.isdir(uacqr_base_path):
    if uacqr_base_path not in sys.path:
        sys.path.insert(0, uacqr_base_path)  # Insert at beginning
        print(f"Added UACQR path: {uacqr_base_path}")
else:
    print(f"Warning: UACQR directory not found at expected location: {uacqr_base_path}")

# Add additional paths that may contain UACQR
sys.path.append(os.path.join(script_dir))
sys.path.append(os.path.join(script_dir, 'UACQR'))

clear_module_path = os.path.join(project_root, "src", "clear")
if clear_module_path not in sys.path:
    sys.path.append(clear_module_path)
    print(f"Appended CLEAR module path: {clear_module_path}")

try:
    from clear.metrics import evaluate_intervals as evaluate_clear_metrics
    from clear.metrics import picp, quantile_loss # For UACQR metrics calculation
    from uacqr_parallel import uacqr  # Now importing directly after path setup
    # Add import for helper module if needed for other functions
    import helper  # Import the helper module for UACQR
except ImportError as e:
    print(f"Error importing modules: {e}")
    print("Please ensure the script is run from the root directory of the 'clear' project")
    print("and both 'src' and 'UACQR' directories are present and contain the necessary files.")
    print("Make sure the UACQR directory contains uacqr_parallel.py and helper.py")
    sys.exit(1)

def load_data_splits(pkl_file_path):
    """Loads data splits from a pickle file generated by benchmark_pcs_data.py.
    Yields X_train, y_train, X_calib, y_calib, X_test, y_test, pcs_median_test, run_key for each run/seed.
    """
    try:
        with open(pkl_file_path, 'rb') as f:
            data_dict = pickle.load(f)

        if not isinstance(data_dict, dict):
            print(f"  Error: Expected a dictionary in {pkl_file_path}, but got {type(data_dict)}")
            # Yield None to indicate failure for this file if it's not a dict of runs
            yield None, None, None, None, None, None, None, None
            return

        for run_key, run_data in data_dict.items():
            print(f"    Processing run/seed: {run_key} from {os.path.basename(pkl_file_path)}")
            if not isinstance(run_data, dict):
                print(f"      Warning: Expected run_data for key '{run_key}' to be a dict, but got {type(run_data)}. Skipping this run.")
                continue

            # Map keys used in benchmark_pcs_data.py to UACQR's expected names
            X_train = np.array(run_data.get("x_train"))
            y_train = np.array(run_data.get("y_train")).flatten()
            X_calib = np.array(run_data.get("x_val")) # Use validation set for calibration
            y_calib = np.array(run_data.get("y_val")).flatten()
            X_test = np.array(run_data.get("x_test"))
            y_test = np.array(run_data.get("y_test")).flatten()

            # Extract PCS median predictions for test data if available
            pcs_median_test = None
            if "test_intervals_raw" in run_data:
                try:
                    # The test_intervals_raw typically contains [lower, median, upper] values
                    test_intervals_raw = np.array(run_data["test_intervals_raw"])
                    pcs_median_test = test_intervals_raw[:, 1]  # Median is in the middle column
                    print(f"      Successfully extracted PCS median predictions for evaluation")
                except (IndexError, ValueError) as e:
                    print(f"      Warning: Could not extract PCS median from test_intervals_raw: {e}")

            # Ensure y arrays are 1D
            if y_train.ndim > 1 and y_train.shape[1] == 1: y_train = y_train.ravel()
            if y_calib.ndim > 1 and y_calib.shape[1] == 1: y_calib = y_calib.ravel()
            if y_test.ndim > 1 and y_test.shape[1] == 1: y_test = y_test.ravel()

            print(f"      Loaded data shapes for run {run_key}: X_train={X_train.shape}, y_train={y_train.shape}, X_calib={X_calib.shape}, y_calib={y_calib.shape}, X_test={X_test.shape}, y_test={y_test.shape}")

            if X_train.shape[0] != y_train.shape[0] or \
               X_calib.shape[0] != y_calib.shape[0] or \
               X_test.shape[0] != y_test.shape[0]:
                print(f"      ValueError: Mismatch in number of samples for run {run_key}. Skipping this run.")
                continue
            
            yield X_train, y_train, X_calib, y_calib, X_test, y_test, pcs_median_test, run_key

    except Exception as e:
        print(f"Error loading or processing data from {pkl_file_path}: {e}")
        # Yield None to indicate failure for this file in case of other exceptions
        yield None, None, None, None, None, None, None, None

def run_uacqr(X_train, y_train, X_calib, y_calib, X_test, alpha, model_params, random_state, current_seed_key):
    """Runs the UACQR method and returns test set predictions and the method name used."""
    try:
        # model_params dict from main includes 'max_features', 'B' (from args)

        # Calculate q_lower and q_upper for the quantile regressors based on the target alpha
        q_lower_val = (alpha / 2) * 100
        q_upper_val = (1 - alpha / 2) * 100
        
        qforest_hyperparams = {
            "min_samples_leaf": 10,  # Fixed to match QRF in train_pcs_quantile.py
            "n_estimators": 100,     # Fixed to match QRF in train_pcs_quantile.py
            "max_features": model_params.get("max_features"), # From args.max_features
        }
        num_bootstraps_uacqr = model_params.get("B") # This is for UACQR's own bootstrapping, from args.n_bootstraps
        uacqr_parallel_jobs = model_params.get("parallel_jobs", -1) # Get parallel_jobs parameter with default of -1

        print(f"    Initializing UACQR. Target alpha for conformalization: {alpha:.4f}, Target quantiles: ({q_lower_val:.2f}, {q_upper_val:.2f}).")
        print(f"    Internal RFQR hyperparams: {qforest_hyperparams}")
        print(f"    UACQR ensemble bootstraps (B): {num_bootstraps_uacqr}")
        print(f"    UACQR bootstrap parallelism: {uacqr_parallel_jobs} jobs")
        
        uacqr_model = uacqr(
            qforest_hyperparams,          # Positional argument for internal RFQR hyperparameters
            model_type='rfqr',            # model_type for the quantile regressor
            alpha=alpha,                  # Alpha for conformal calibration
            q_lower=q_lower_val,          # Lower quantile for base regressors
            q_upper=q_upper_val,          # Upper quantile for base regressors
            B=num_bootstraps_uacqr,       # Number of UACQR bootstrap models
            random_state=random_state,    # Pass main random_state for UACQR's own processes
            uacqrs_agg='iqr',             # Aggregation method for UACQR-S, as in demo_rf.py
            double_bootstrapping=True,    # Enable internal bootstrapping for the B models
            n_jobs=uacqr_parallel_jobs    # Set number of parallel jobs for bootstrap model fitting
        )

        start_time = time.time()

        print("    Fitting UACQR base models...")
        uacqr_model.fit(X_train, y_train)

        print("    Calibrating UACQR models...")
        uacqr_model.calibrate(X_calib, y_calib)

        print("    Predicting UACQR intervals for test set...")
        # UACQR's predict method returns a pandas DataFrame with all methods
        predictions_df = uacqr_model.predict(X_test)
        
        print(f"    UACQR predict() returned DataFrame with columns: {list(predictions_df.columns.levels[0]) if isinstance(predictions_df.columns, pd.MultiIndex) else list(predictions_df.columns)}")

        end_time = time.time()
        fit_duration = end_time - start_time
        print(f"    UACQR model fitting took {fit_duration:.4f} seconds.")

        # Return the entire DataFrame of predictions
        return predictions_df

    except Exception as e:
        print(f"Error running UACQR: {e}")
        import traceback
        traceback.print_exc()
        return None

def evaluate_performance(y_test, lower, upper, f_pred, alpha):
    """Evaluates intervals using metrics from src.clear.metrics.py."""
    results = {}

    # --- Calculate metrics using evaluate_intervals (aliased as evaluate_clear_metrics) ---
    try:
        print("      Calculating metrics using evaluate_intervals from src.clear.metrics.py...")
        # Directly use the metric names as returned by evaluate_intervals
        metrics_from_clear_eval = evaluate_clear_metrics(y_test, lower, upper, alpha=alpha, f=f_pred)
        results.update(metrics_from_clear_eval) # Add all calculated metrics
    except Exception as e:
        print(f"      Error calculating metrics with evaluate_clear_metrics: {e}")
        # Initialize with NaN if calculation fails to maintain consistent columns
        metric_keys = ["PICP", "NIW", "MPIW", "QuantileLoss", "CRPS", "AUC", "NCIW", "c_test_cal", "ExpectileLoss", "IntervalScoreLoss"]
        for key in metric_keys:
            results[key] = np.nan
    
    # The UACQR-specific metrics (Coverage, AvgLength, IntervalScore) are now effectively
    # covered by PICP, MPIW, and QuantileLoss from evaluate_intervals.
    # So, the separate UACQR metric calculation block is removed.

    return results

def load_clear_results(dataset_name_base, seed_key_str, clear_results_dir, target_alpha_eval):
    """Loads 'clear' results for a given dataset and seed 
       from a pre-existing benchmark CSV file.

    Args:
        dataset_name_base (str): The base name of the dataset (e.g., 'ailerons').
        seed_key_str (str): The string identifying the seed/run (e.g., 'run_0').
        clear_results_dir (str): Directory containing the CLEAR benchmark CSVs.
        target_alpha_eval (float): The alpha value used for UACQR evaluation, for consistency.

    Returns:
        list: A list of result_row dictionaries for the found CLEAR methods, or an empty list.
    """
    
    # Infer confidence from target_alpha_eval to match filename, e.g. 0.05 -> 95
    confidence_level_int = int((1 - target_alpha_eval) * 100)
    csv_filename = f"benchmark_results_data_{dataset_name_base}_{confidence_level_int}.csv"
    csv_filepath = os.path.join(clear_results_dir, csv_filename)

    if not os.path.exists(csv_filepath):
        print(f"      INFO: CLEAR results CSV not found for {dataset_name_base} at {csv_filepath}. Skipping CLEAR data loading for this dataset.")
        return []

    try:
        df_clear = pd.read_csv(csv_filepath)
    except Exception as e:
        print(f"      ERROR: Could not read CLEAR results CSV {csv_filepath}: {e}")
        return []

    # Filter for the specific seed and relevant methods
    try:
        seed_val = int(seed_key_str.split('_')[-1])
    except ValueError:
        print(f"      Warning: Could not parse seed integer from seed_key '{seed_key_str}'. Skipping CLEAR data for this seed.")
        return []

    relevant_methods = ['clear'] # Only load 'clear'
    df_filtered = df_clear[
        (df_clear['Dataset'] == f"data_{dataset_name_base}") &
        (df_clear['Method'].isin(relevant_methods)) &
        (df_clear['Seed'] == seed_val)
    ]

    if df_filtered.empty:
        print(f"      INFO: No 'clear' methods found for dataset {dataset_name_base}, seed {seed_val} in {csv_filepath}.")
        return []

    loaded_results = []
    for method_name in relevant_methods:
        df_method_seed = df_filtered[df_filtered['Method'] == method_name]
        if df_method_seed.empty:
            continue
        
        metric_dict = {}
        # Pivot the table to get metrics as columns
        # Metric names in CSV are lowercase, in benchmark_uacqr.py they are uppercase
        for _, row in df_method_seed.iterrows():
            metric_key_csv = str(row['Metric']).lower() # e.g., 'picp', 'niw'
            # Map to uppercase keys used in benchmark_uacqr.py cols_order
            # PICP, NIW, MPIW, QuantileLoss, CRPS, AUC, NCIW, c_test_cal
            if metric_key_csv == 'picp': metric_dict['PICP'] = row['Value']
            elif metric_key_csv == 'niw': metric_dict['NIW'] = row['Value']
            elif metric_key_csv == 'mpiw': metric_dict['MPIW'] = row['Value']
            elif metric_key_csv == 'quantileloss': metric_dict['QuantileLoss'] = row['Value']
            elif metric_key_csv == 'crps': metric_dict['CRPS'] = row['Value']
            elif metric_key_csv == 'auc': metric_dict['AUC'] = row['Value']
            elif metric_key_csv == 'nciw': metric_dict['NCIW'] = row['Value']
            elif metric_key_csv == 'c_test_cal': metric_dict['c_test_cal'] = row['Value']
            elif metric_key_csv == 'expectileloss': metric_dict['ExpectileLoss'] = row['Value']
            elif metric_key_csv == 'intervalscoreloss': metric_dict['IntervalScoreLoss'] = row['Value']
            # Add other metrics if needed, ensuring consistent naming

        if not metric_dict: # If no relevant metrics were found for this method/seed
            continue
            
        result_row_clear = {
            "Dataset": dataset_name_base,
            "Seed": seed_key_str, # Use the original string key
            "Method": method_name, # 'clear' or 'clear_c'
            "Coverage_Target": (1 - target_alpha_eval),
            "Alpha": target_alpha_eval,
            **metric_dict
        }
        loaded_results.append(result_row_clear)
        print(f"      Successfully loaded {method_name} results for seed {seed_key_str} from {csv_filename}")

    return loaded_results

def main(args):
    """Main execution function."""
    # --- Logging Setup ---
    log_level_arg = getattr(logging, args.log_level.upper(), logging.INFO)
    log_file_arg = None

    if args.log_file.lower() != "none":
        if args.log_file.lower() == "auto":
            # Create a 'logs' directory at the project root if it doesn't exist
            log_dir_path = os.path.join(project_root, "logs")
            # os.makedirs(log_dir_path, exist_ok=True) # setup_logging will handle this
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            log_file_arg = os.path.join(log_dir_path, f"benchmark_uacqr_{timestamp}.log")
        else:
            # User specified a log file path
            # If it's not absolute, make it relative to project_root
            if not os.path.isabs(args.log_file):
                log_file_arg = os.path.join(project_root, args.log_file)
            else:
                log_file_arg = args.log_file
    
    logger = setup_logging(log_file_arg, log_level_arg)
    sys.stdout = StreamToLogger(logger, logging.INFO) # Redirect print to logger.info
    # --- End Logging Setup ---

    logger.info("Starting UACQR Benchmark Script")
    logger.info(f"Script arguments: {args}")
    logger.info(f"Project root: {project_root}")
    # Potentially log other relevant versions or environment info

    # Convert relative paths to absolute paths
    
    # --- Handle data_dir ---
    original_data_dir = args.data_dir
    if not os.path.isabs(args.data_dir):
        if args.data_dir.startswith(".."):
            # Relative to script dir
            args.data_dir = os.path.abspath(os.path.join(script_dir, args.data_dir))
            logger.info(f"Converting relative data_dir path '{original_data_dir}' (from script location) to absolute: '{args.data_dir}'")
        else:
            # Relative to project root
            args.data_dir = os.path.abspath(os.path.join(project_root, args.data_dir))
            logger.info(f"Converting relative data_dir path '{original_data_dir}' (from project root) to absolute: '{args.data_dir}'")
    else:
        logger.info(f"Data_dir path '{original_data_dir}' is absolute. Using as is.")

    # Find dataset files
    pkl_pattern = os.path.join(args.data_dir, f"data_*.pkl")
    pkl_files = glob.glob(pkl_pattern)

    if not pkl_files:
        logger.error(f"Error: No data pickle files found in {args.data_dir} matching 'data_*.pkl'")
        sys.exit(1)

    # --- Handle output_csv ---
    original_output_csv = args.output_csv
    if not os.path.isabs(args.output_csv):
        if args.output_csv.startswith(".."):
            # Relative to script dir
            args.output_csv = os.path.abspath(os.path.join(script_dir, args.output_csv))
            logger.info(f"Converting relative output_csv path '{original_output_csv}' (from script location) to absolute: '{args.output_csv}'")
        else:
            # Relative to project root
            args.output_csv = os.path.abspath(os.path.join(project_root, args.output_csv))
            logger.info(f"Converting relative output_csv path '{original_output_csv}' (from project root) to absolute: '{args.output_csv}'")
    else:
        logger.info(f"Output_csv path '{original_output_csv}' is absolute. Using as is.")
        
    # --- Handle clear_results_dir ---
    if args.clear_results_dir:
        original_clear_results_dir = args.clear_results_dir
        if not os.path.isabs(args.clear_results_dir):
            if args.clear_results_dir.startswith(".."):
                # Relative to script dir
                args.clear_results_dir = os.path.abspath(os.path.join(script_dir, args.clear_results_dir))
                logger.info(f"Converting relative clear_results_dir path '{original_clear_results_dir}' (from script location) to absolute: '{args.clear_results_dir}'")
            else:
                # Relative to project root
                args.clear_results_dir = os.path.abspath(os.path.join(project_root, args.clear_results_dir))
                logger.info(f"Converting relative clear_results_dir path '{original_clear_results_dir}' (from project root) to absolute: '{args.clear_results_dir}'")
        else:
            logger.info(f"Clear_results_dir path '{original_clear_results_dir}' is absolute. Using as is.")

    logger.info(f"Found {len(pkl_files)} dataset files in {args.data_dir}")

    all_results = []
    alpha = 1.0 - args.coverage

    # Column order for CSV output (defined once, used for incremental + final writes)
    _cols_order = ["Dataset", "Seed", "Method", "Coverage_Target", "Alpha",
                   "PICP", "NIW", "MPIW", "QuantileLoss", "ExpectileLoss", "IntervalScoreLoss",
                   "CRPS", "AUC", "NCIW", "c_test_cal"]

    # Prepare output CSV for incremental writing (clear any previous file)
    _output_dir = os.path.dirname(args.output_csv)
    if _output_dir and not os.path.exists(_output_dir):
        os.makedirs(_output_dir)
    if os.path.exists(args.output_csv):
        os.remove(args.output_csv)
        logger.info(f"Cleared previous output file: {args.output_csv}")
    _written_idx = 0

    # Define UACQR parameters from args
    uacqr_params = {
        # "min_samples_leaf": args.min_samples_leaf, # Not used for internal RFQR hyperparams anymore
        "max_features": args.max_features, # Still used for internal RFQR
        "B": args.n_bootstraps, # Pass bootstrap number for UACQR ensemble
        "parallel_jobs": args.parallel_jobs # Pass the number of parallel jobs
    }

    # Loop through datasets
    for pkl_file in tqdm(pkl_files, desc="Processing datasets"):
        dataset_name_full = os.path.basename(pkl_file).replace('data_', '').replace('.pkl', '')
        # Extract base dataset name if it includes seed info like "_run_X" or "_pcs_results_95"
        dataset_name_base = dataset_name_full.split('_pcs_results')[0].split('_run_')[0]
        logger.info(f"\n--- Processing dataset file: {os.path.basename(pkl_file)} (Base: {dataset_name_base}) ---")

        # Load data - now a generator
        for X_train, y_train, X_calib, y_calib, X_test, y_test, pcs_median_test, seed_key in load_data_splits(pkl_file):
            if X_train is None: # Check if loading failed for a run/seed
                logger.warning(f"  Skipping a run/seed in {dataset_name_full} due to loading error for that specific run.")
                continue

            current_run_random_state = args.random_state + int(seed_key.split('_')[-1]) if 'run_' in seed_key else args.random_state

            logger.info(f"  --- Running UACQR for seed: {seed_key} (Random state for UACQR models: {current_run_random_state}) ---")
            # Run UACQR - it now returns the full predictions DataFrame
            all_uacqr_predictions_df = run_uacqr(
                X_train, y_train, X_calib, y_calib, X_test,
                alpha=alpha, # This is 1.0 - args.coverage
                model_params=uacqr_params,
                random_state=current_run_random_state,
                current_seed_key=seed_key
            )

            if all_uacqr_predictions_df is None:
                logger.warning(f"  Skipping evaluation for {dataset_name_base} (seed {seed_key}) due to UACQR error.")
                continue

            # Methods from UACQR to evaluate and log
            uacqr_methods_to_log = ['UACQR-S', 'UACQR-P'] 

            for method_name_uacqr in uacqr_methods_to_log:
                if method_name_uacqr not in all_uacqr_predictions_df.columns.levels[0]:
                    logger.warning(f"    Warning: Method '{method_name_uacqr}' not found in UACQR predictions output. Skipping for this method.")
                    continue

                logger.info(f"    Extracting and evaluating {method_name_uacqr} predictions for seed {seed_key}...")
                uacqr_lower = all_uacqr_predictions_df[method_name_uacqr]['lower'].values
                uacqr_upper = all_uacqr_predictions_df[method_name_uacqr]['upper'].values
                
                # Use the PCS median for evaluation if available, otherwise use midpoint
                if pcs_median_test is not None:
                    uacqr_f = pcs_median_test  # Use PCS median from original data
                    logger.info(f"      Using PCS median for {method_name_uacqr} evaluation")
                else:
                    # Fall back to midpoint if PCS median not available
                    uacqr_f = (uacqr_lower + uacqr_upper) / 2.0
                    logger.info(f"      PCS median not available, using midpoint for {method_name_uacqr} evaluation")
                
                logger.info(f"      Generated {method_name_uacqr} predictions: lower={uacqr_lower.shape}, upper={uacqr_upper.shape}")

                # Evaluate results for the current UACQR method
                uacqr_metrics = evaluate_performance(y_test, uacqr_lower, uacqr_upper, uacqr_f, alpha)

                # Store UACQR results for the current method
                result_row_uacqr = {
                    "Dataset": dataset_name_base,
                    "Seed": seed_key, 
                    "Method": method_name_uacqr, # Store specific method name (UACQR-S or UACQR-P)
                    "Coverage_Target": args.coverage,
                    "Alpha": alpha,
                    **uacqr_metrics 
                }
                all_results.append(result_row_uacqr)
                logger.info(f"    Results for {dataset_name_base} - {method_name_uacqr} (seed {seed_key}):")
                for metric, value in uacqr_metrics.items():
                    logger.info(f"      {metric}: {value}")

            # Load and print CLEAR results for this dataset and seed if path provided
            if args.clear_results_dir:
                clear_data_list = load_clear_results(dataset_name_base, seed_key, args.clear_results_dir, alpha)
                if clear_data_list:
                    all_results.extend(clear_data_list)
                    for clear_result_row in clear_data_list: # Should be only one if found
                        logger.info(f"    Results for {dataset_name_base} - {clear_result_row['Method']} (seed {seed_key}):")
                        for metric, value in clear_result_row.items():
                            # Skip printing auxiliary info already covered or not metrics
                            if metric not in ["Dataset", "Seed", "Method", "Coverage_Target", "Alpha"]:
                                logger.info(f"      {metric}: {value}")

            # For verification
            if seed_key == 'run_0':  # Only for the first run to avoid too much output
                for method in ["UACQR-P", "clear"]:
                    # Calculate just the miss rate
                    if method == "UACQR-P":
                        miss_rate = 1.0 - picp(y_test, uacqr_lower, uacqr_upper)
                        # Calculate average miss magnitude
                        misses_below = (y_test < uacqr_lower)
                        misses_above = (y_test > uacqr_upper)
                        avg_miss_mag = np.mean(np.concatenate([
                            uacqr_lower[misses_below] - y_test[misses_below],
                            y_test[misses_above] - uacqr_upper[misses_above]
                        ]))
                        logger.info(f"DEBUG - {method} miss rate: {miss_rate:.4f}, avg miss magnitude: {avg_miss_mag:.6f}")

        # Incremental save after each dataset completes
        if len(all_results) > _written_idx:
            _new_df = pd.DataFrame(all_results[_written_idx:])
            _cols = [c for c in _cols_order if c in _new_df.columns]
            _new_df = _new_df[_cols]
            _header = not os.path.exists(args.output_csv)
            _new_df.to_csv(args.output_csv, mode='a', index=False,
                           header=_header, float_format='%.10g')
            _written_idx = len(all_results)
            logger.info(f"  [Incremental save] {dataset_name_base} done — {_written_idx} rows total in {args.output_csv}")

    # Final flush: write any rows not yet saved incrementally
    if not all_results:
        logger.info("\nNo results generated.")
        return

    if _written_idx < len(all_results):
        _new_df = pd.DataFrame(all_results[_written_idx:])
        _cols = [c for c in _cols_order if c in _new_df.columns]
        _new_df = _new_df[_cols]
        _header = not os.path.exists(args.output_csv)
        _new_df.to_csv(args.output_csv, mode='a', index=False,
                       header=_header, float_format='%.10g')

    logger.info(f"\nBenchmark complete. {len(all_results)} total rows saved to: {args.output_csv}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Benchmark UACQR method on pre-processed datasets.")
    parser.add_argument("--data_dir", type=str, default="../../models/pcs_top1_pcs_10_standard",
                        help="Directory containing the dataset .pkl files (e.g., '../../models/pcs_top1_pcs_10_standard').")
    parser.add_argument("--output_csv", type=str, default="../../results/uacqr/uacqr_benchmark_results_standard.csv",
                        help="Path to save the benchmark results CSV file.")
    parser.add_argument("--coverage", type=float, default=0.95,
                        help="Desired prediction interval coverage level (e.g., 0.90, 0.95).")
    parser.add_argument("--n_bootstraps", type=int, default=100,
                        help="Number of bootstrap models for UACQR (B parameter for UACQR ensemble).")
    parser.add_argument("--min_samples_leaf", type=int, default=10, # Changed default to 10, though it's fixed to 10 in run_uacqr for internal RFQR
                        help="min_samples_leaf for UACQR's internal RandomForestQuantileRegressor (fixed to 10 in code). This arg is currently not directly used for that.")
    parser.add_argument("--max_features", type=float, default=1.0, # Added back as it IS used for qforest_hyperparams
                         help="max_features parameter for UACQR's internal RandomForestQuantileRegressor (as fraction).")
    parser.add_argument("--random_state", type=int, default=777,
                        help="Base random seed for reproducibility. Seed for each run will be base_seed + run_index.")
    parser.add_argument("--clear_results_dir", type=str, default='../../results/standard/PCS_all_10seeds_qrf',
                        help="Directory containing the CLEAR benchmark CSV files (e.g., '../../results/PCS_all_10seeds_qrf') to load comparative results from.")
    parser.add_argument("--parallel_jobs", type=int, default=30,
                        help="Number of parallel jobs for UACQR's bootstrap model fitting. Use -1 for all available cores.")
    parser.add_argument("--log_file", type=str, default="auto",
                        help="Path to log file. Use 'none' to disable file logging. Default 'auto' creates timestamped logs in a 'logs' directory relative to project root.")
    parser.add_argument("--log_level", type=str, default="INFO", 
                        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
                        help="Logging level.")

    args = parser.parse_args()

    # Validate coverage
    if not 0 < args.coverage < 1:
        # Use logger for error messages before exiting
        # print to stderr as a fallback, then try to log if possible.
        eprint = lambda *x, **kw: print(*x, file=sys.stderr, **kw)
        eprint("Error: Coverage must be between 0 and 1.")
        try:
            logging.getLogger().error("Error: Coverage must be between 0 and 1.")
        except: # Fallback if logger not set up
            pass
        sys.exit(1)

    main(args) 