"""
Main Orchestrator for the Trading Pipeline.

Coordinates data loading, feature engineering, model training, calibration, 
SAC training, and backtesting.
"""

import os
import sys
import logging
import yaml
import pandas as pd
import numpy as np
from datetime import datetime
import argparse
import joblib
import json

# Determine the project root directory based on the script location
# This assumes the script is in src/ and the project root is two levels up
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(script_dir))
# Add project root to sys.path to allow absolute imports from the package
if project_root not in sys.path:
    sys.path.insert(0, project_root)

# Now use absolute imports based on the package structure
from gru_sac_predictor.src.data_loader import DataLoader
# Import other components as they are created
from gru_sac_predictor.src.feature_engineer import FeatureEngineer
# Try importing minimal_whitelist from features.py (assuming it exists there)
try:
    from gru_sac_predictor.src.features import minimal_whitelist
except ImportError:
    # Fallback: Define it here if features.py doesn't exist or doesn't have it
    logging.warning("Could not import minimal_whitelist from .features, defining fallback.")
    minimal_whitelist = [
        "return_1m", "return_15m", "return_60m", "ATR_14", "volatility_14d",
        "chaikin_AD_10", "svi_10", "EMA_10", "EMA_50", "MACD", "MACD_signal",
        "hour_sin", "hour_cos",
    ]
from gru_sac_predictor.src.gru_model_handler import GRUModelHandler
from gru_sac_predictor.src.calibrator import Calibrator
from gru_sac_predictor.src.sac_trainer import SACTrainer
from gru_sac_predictor.src.backtester import Backtester

# Removed redundant imports for feature selection
from sklearn.preprocessing import StandardScaler


logger = logging.getLogger(__name__) # Use module-level logger

class TradingPipeline:
    """Orchestrates the entire trading strategy pipeline."""

    def __init__(self, config_path: str):
        """Initialize the pipeline with configuration."""
        self.config_path = config_path
        self.config = self._load_config()
        self.run_id = self._generate_run_id()
        self._setup_directories()
        self._setup_logging()
        logging.info(f"--- Starting Pipeline Run: {self.run_id} ---")
        logging.info(f"Using config: {self.config_path}")

        # Instantiate components
        db_dir_from_config = self.config['data']['db_dir']
        self.data_loader = DataLoader(db_dir=db_dir_from_config)
        self.feature_engineer = FeatureEngineer(minimal_whitelist=minimal_whitelist)
        self.gru_handler = GRUModelHandler(run_id=self.run_id, models_dir=self.current_run_models_dir)
        cal_cfg = self.config.get('calibration', {})
        self.calibrator = Calibrator(edge_threshold=cal_cfg.get('edge_threshold', 0.55))
        self.sac_trainer = None
        self.backtester = Backtester(config=self.config)

        # Initialize data/state variables
        self.df_raw = None
        self.df_engineered_full = None
        self.df_features_minimal = None
        self.df_targets = None
        self.df_train = None
        self.df_val = None
        self.df_test = None
        self.X_train_raw = None
        self.X_val_raw = None
        self.X_test_raw = None
        self.y_train = None
        self.y_val = None
        self.y_test = None
        self.y_dir_train = None
        self.final_whitelist = None
        self.scaler = None
        self.X_train_pruned = None
        self.X_val_pruned = None
        self.X_test_pruned = None
        self.X_train_scaled = None
        self.X_val_scaled = None
        self.X_test_scaled = None
        self.X_train_seq, self.y_train_seq_dict = None, None
        self.X_val_seq, self.y_val_seq_dict = None, None
        self.X_test_seq, self.y_test_seq_dict = None, None
        self.gru_model = None
        self.gru_model_run_id_loaded_from = None
        self.optimal_T = None
        self.sac_agent_load_path = None
        self.train_indices, self.val_indices, self.test_indices = None, None, None
        self.backtest_results_df = None
        self.backtest_metrics = None

        self._save_run_config()

    def _load_config(self) -> dict:
        """Loads the YAML configuration file."""
        try:
            # Try loading relative to the script first (if running from src)
            if not os.path.isabs(self.config_path):
                potential_path = os.path.join(script_dir, self.config_path)
                if not os.path.exists(potential_path):
                    # If not found relative to script, try relative to project root
                    potential_path = os.path.join(project_root, self.config_path)
                    if not os.path.exists(potential_path):
                       # If still not found, try relative to CWD as last resort
                       potential_path = os.path.abspath(self.config_path)
                
                if os.path.exists(potential_path):
                     self.config_path = potential_path
                else:
                    # Try one level up from project root (common structure)
                    potential_path = os.path.join(os.path.dirname(project_root), 'gru_sac_predictor', 'config.yaml')
                    if os.path.exists(potential_path):
                         self.config_path = potential_path
                    else:
                        raise FileNotFoundError(f"Config file not found at relative paths, CWD, or common location: {self.config_path}")
            
            with open(self.config_path, 'r') as f:
                config = yaml.safe_load(f)
            # Basic validation
            if 'data' not in config or 'gru' not in config or 'sac' not in config:
                raise ValueError("Config file missing essential sections: data, gru, sac")
            # Validate calibration config if present
            if 'calibration' in config and 'edge_threshold' not in config['calibration']:
                 logging.warning("'edge_threshold' not found in calibration config, using default 0.55")
                 config['calibration']['edge_threshold'] = 0.55 # Add default if missing
            elif 'calibration' not in config:
                 logging.warning("'calibration' section not found in config, using default edge_threshold 0.55")
                 config['calibration'] = {'edge_threshold': 0.55} # Add default section
                 
            return config
        except FileNotFoundError:
            print(f"ERROR: Configuration file not found at '{self.config_path}'")
            sys.exit(1)
        except yaml.YAMLError as e:
            print(f"ERROR: Error parsing configuration file '{self.config_path}': {e}")
            sys.exit(1)
        except Exception as e:
            print(f"ERROR: An unexpected error occurred while loading config: {e}")
            sys.exit(1)

    def _generate_run_id(self) -> str:
        """Generates a unique run ID based on the template in config."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        template = self.config.get('run_id_template', '{timestamp}')
        return template.format(timestamp=timestamp)

    def _setup_directories(self):
        """Creates directories for logs, results, and models for this run."""
        self.dirs = {}
        base_dirs_config = self.config.get('base_dirs', {})
        # Calculate base models dir path (needed for loading previous models)
        # Assume it's relative to project root
        models_rel_path = base_dirs_config.get('models', 'models')
        self.base_models_dir_path = os.path.join(project_root, models_rel_path)
        
        for dir_type, rel_path in base_dirs_config.items():
            # Paths are relative to the project root
            abs_path = os.path.join(project_root, rel_path, self.run_id)
            os.makedirs(abs_path, exist_ok=True)
            self.dirs[dir_type] = abs_path
            # No need to log here, happens in _setup_logging
        
        # Specific dir for current run models (if models base dir exists)
        if 'models' in self.dirs:
            self.current_run_models_dir = self.dirs['models']
        else:
            # Fallback if models base dir not in config
            self.current_run_models_dir = os.path.join(self.base_models_dir_path, self.run_id) # Use calculated base path
            os.makedirs(self.current_run_models_dir, exist_ok=True)
            # Log this warning after logging is set up
            # logging.warning(f"'models' base dir not found in config, using default: {self.current_run_models_dir}")

    def _setup_logging(self):
        """Configures logging to file and console."""
        log_dir = self.dirs.get('logs')
        if not log_dir:
             print(f"Warning: 'logs' directory not configured. Logging to console only.")
             log_file_path = None
        else:
             log_file_path = os.path.join(log_dir, f'pipeline_{self.run_id}.log')

        # Remove existing handlers to avoid duplicate logs if re-initialized
        for handler in logging.root.handlers[:]:
            logging.root.removeHandler(handler)

        log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        log_level = logging.INFO # Consider making this configurable

        handlers = [logging.StreamHandler(sys.stdout)]
        if log_file_path:
             handlers.append(logging.FileHandler(log_file_path))

        logging.basicConfig(level=log_level, format=log_format, handlers=handlers)

        # Configure TensorFlow logging (optional, reduces verbosity)
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # ERROR messages only
        logging.getLogger('tensorflow').setLevel(logging.ERROR)
        
        # Now log the directory paths
        logging.info(f"Using Base Models Directory: {self.base_models_dir_path}")
        for dir_type, abs_path in self.dirs.items():
             logging.info(f"Using {dir_type} directory: {abs_path}")
        if 'models' not in self.dirs:
             logging.warning(f"'models' base dir not found in config, using default models dir: {self.current_run_models_dir}")

        logging.info(f"Logging setup complete. Log file: {log_file_path if log_file_path else 'Console only'}")

    def _save_run_config(self):
        """Saves the configuration used for this run."""
        results_dir = self.dirs.get('results')
        if results_dir:
            config_save_path = os.path.join(results_dir, 'run_config.yaml')
            try:
                with open(config_save_path, 'w') as f:
                    yaml.dump(self.config, f, default_flow_style=False)
                logging.info(f"Saved run configuration to {config_save_path}")
            except Exception as e:
                logging.error(f"Failed to save run configuration: {e}")
        else:
            logging.warning("'results' directory not configured. Skipping saving run config.")

    # --- Pipeline Stages ---

    def load_and_preprocess_data(self):
        """Loads raw data using DataLoader and performs initial checks."""
        logging.info("--- Stage: Loading and Preprocessing Data ---")
        data_cfg = self.config['data']
        self.df_raw = self.data_loader.load_data(
            ticker=data_cfg['ticker'],
            exchange=data_cfg['exchange'],
            start_date=data_cfg['start_date'],
            end_date=data_cfg['end_date'],
            interval=data_cfg['interval']
        )
        if self.df_raw is None or self.df_raw.empty:
            logging.error("Failed to load data. Exiting.")
            sys.exit(1)

        # Basic checks
        if not isinstance(self.df_raw.index, pd.DatetimeIndex):
            logging.error("Data index is not a DatetimeIndex after loading. Exiting.")
            sys.exit(1)
        if self.df_raw.index.tz is None or self.df_raw.index.tz.zone.upper() != 'UTC': # Case-insensitive check
            logging.warning(f"Data index timezone is not UTC ({self.df_raw.index.tz}). Attempting conversion.")
            try:
                if self.df_raw.index.tz is None:
                    self.df_raw = self.df_raw.tz_localize('UTC')
                else:
                    self.df_raw = self.df_raw.tz_convert('UTC')
                logging.info(f"Data index timezone converted to UTC.")
            except Exception as e:
                 logging.error(f"Failed to convert index timezone to UTC: {e}. Exiting.")
                 sys.exit(1)
                 
        # Drop rows with NaN in essential OHLCV columns early
        initial_rows = len(self.df_raw)
        self.df_raw.dropna(subset=['open', 'high', 'low', 'close', 'volume'], inplace=True)
        if len(self.df_raw) < initial_rows:
            logging.warning(f"Dropped {initial_rows - len(self.df_raw)} rows with NaN in OHLCV during loading.")

        logging.info(f"Raw data loaded successfully: {self.df_raw.shape[0]} rows from {self.df_raw.index.min()} to {self.df_raw.index.max()}")

    def engineer_features(self):
        """Adds features using FeatureEngineer."""
        logging.info("--- Stage: Engineering Features ---")
        if self.df_raw is None:
            logging.error("Raw data not loaded. Cannot engineer features.")
            sys.exit(1)
        
        # Add base features (cyclical, imbalance, TA)
        self.df_engineered_full = self.feature_engineer.add_base_features(self.df_raw.copy())

        # Prune to minimal whitelist immediately for comparison/logging if needed
        # self.df_features_minimal = self.feature_engineer.prune_features(self.df_engineered_full, self.feature_engineer.minimal_whitelist)
        # logger.info(f"Features pruned to minimal whitelist ({self.df_features_minimal.shape[1]}): {self.df_features_minimal.columns.tolist()}")

        # Drop rows with NaNs potentially introduced by feature engineering (especially rolling features at the start)
        initial_rows = len(self.df_engineered_full)
        self.df_engineered_full.dropna(inplace=True)
        if len(self.df_engineered_full) < initial_rows:
             logging.warning(f"Dropped {initial_rows - len(self.df_engineered_full)} rows with NaN values after feature engineering.")

        if self.df_engineered_full.empty:
             logging.error("DataFrame is empty after feature engineering and NaN removal. Exiting.")
             sys.exit(1)

        logging.info(f"Feature engineering complete. Shape: {self.df_engineered_full.shape}")

    def define_labels_and_align(self):
        """Defines prediction labels (returns, direction) and aligns with features."""
        logging.info("--- Stage: Defining Labels and Aligning ---")
        if self.df_engineered_full is None:
            logging.error("Engineered data not available. Cannot define labels.")
            sys.exit(1)

        # Calculate forward returns and direction based on 'close' price
        # Note: Ensure 'close' is present from the original raw data load
        if 'close' not in self.df_engineered_full.columns:
            logging.error("'close' column missing in engineered data. Cannot define labels.")
            sys.exit(1)

        horizon = self.config['gru'].get('prediction_horizon', 5) # Default horizon if not in config
        target_ret_col = f'fwd_log_ret_{horizon}'
        target_dir_col = f'direction_label_{horizon}'

        # Shift close price into the future by 'horizon' periods
        shifted_close = self.df_engineered_full['close'].shift(-horizon)
        
        # Calculate log return
        self.df_engineered_full[target_ret_col] = np.log(shifted_close / self.df_engineered_full['close'])
        
        # Calculate direction label (1 if future price > current price, 0 otherwise)
        self.df_engineered_full[target_dir_col] = (self.df_engineered_full[target_ret_col] > 0).astype(int)

        # Drop rows where targets are NaN (due to the shift at the end of the DataFrame)
        initial_rows = len(self.df_engineered_full)
        self.df_engineered_full.dropna(subset=[target_ret_col, target_dir_col], inplace=True)
        final_rows = len(self.df_engineered_full)
        if final_rows < initial_rows:
            logging.info(f"Dropped {initial_rows - final_rows} rows due to NaN targets (horizon={horizon}).")

        if self.df_engineered_full.empty:
             logging.error("DataFrame is empty after defining labels and dropping NaNs. Exiting.")
             sys.exit(1)
             
        # Separate features (X) and targets (y) - X contains all engineered features for now
        self.X_raw_aligned = self.df_engineered_full.drop(columns=[target_ret_col, target_dir_col])
        self.y_aligned = self.df_engineered_full[[target_ret_col, target_dir_col]]
        self.y_dir_aligned = self.df_engineered_full[target_dir_col] # Keep separate handle for feature selection

        logging.info(f"Labels (horizon={horizon}) defined and aligned. Features shape: {self.X_raw_aligned.shape}, Targets shape: {self.y_aligned.shape}")

    def split_data(self):
        """Splits features and targets into train, validation, and test sets chronologically."""
        logging.info("--- Stage: Splitting Data ---")
        if self.X_raw_aligned is None or self.y_aligned is None:
            logging.error("Aligned features/targets not available for splitting.")
            sys.exit(1)
        if not isinstance(self.X_raw_aligned.index, pd.DatetimeIndex):
             logging.error("Feature index must be DatetimeIndex for chronological split. Aborting.")
             sys.exit(1)

        split_cfg = self.config['split_ratios']
        train_ratio = split_cfg['train']
        val_ratio = split_cfg['validation']
        test_ratio = round(1.0 - train_ratio - val_ratio, 2)
        logger.info(f"Using split ratios: Train={train_ratio:.2f}, Val={val_ratio:.2f}, Test={test_ratio:.2f}")

        total_len = len(self.X_raw_aligned)
        train_end_idx = int(total_len * train_ratio)
        val_end_idx = int(total_len * (train_ratio + val_ratio))

        # Split features
        self.X_train_raw = self.X_raw_aligned.iloc[:train_end_idx]
        self.X_val_raw = self.X_raw_aligned.iloc[train_end_idx:val_end_idx]
        self.X_test_raw = self.X_raw_aligned.iloc[val_end_idx:]

        # Split targets
        self.y_train = self.y_aligned.iloc[:train_end_idx]
        self.y_val = self.y_aligned.iloc[train_end_idx:val_end_idx]
        self.y_test = self.y_aligned.iloc[val_end_idx:]
        
        # Keep separate handle to direction target for training feature selector
        self.y_dir_train = self.y_dir_aligned.iloc[:train_end_idx]

        logging.info(f"Data split complete:")
        logging.info(f"  Train: X={self.X_train_raw.shape}, y={self.y_train.shape} ({self.X_train_raw.index.min()} to {self.X_train_raw.index.max()})")
        logging.info(f"  Val:   X={self.X_val_raw.shape}, y={self.y_val.shape} ({self.X_val_raw.index.min()} to {self.X_val_raw.index.max()})")
        logging.info(f"  Test:  X={self.X_test_raw.shape}, y={self.y_test.shape} ({self.X_test_raw.index.min()} to {self.X_test_raw.index.max()})")

        if len(self.X_train_raw) == 0 or len(self.X_val_raw) == 0 or len(self.X_test_raw) == 0:
            logging.error("One or more data splits are empty. Check data length and split ratios. Aborting.")
            sys.exit(1)

    def select_and_prune_features(self):
        """Performs feature selection (e.g., VIF, L1) and prunes data splits."""
        logging.info("--- Stage: Selecting and Pruning Features ---")
        if self.X_train_raw is None or self.y_dir_train is None:
             logging.error("Training data (X_train_raw, y_dir_train) not available for feature selection.")
             sys.exit(1)
        
        # Perform feature selection using the training set
        self.final_whitelist = self.feature_engineer.select_features(
            self.X_train_raw, 
            self.y_dir_train,
            # Optionally get VIF threshold from config? Defaulting for now.
            # vif_threshold=self.config.get('feature_selection', {}).get('vif_threshold', 10.0) 
        )

        # --- Save the final whitelist --- # Should this be done here or in FeatureEngineer?
        # Let's keep it here for pipeline-level artifact saving.
        whitelist_save_path = os.path.join(self.current_run_models_dir, f'final_whitelist_{self.run_id}.json')
        try:
            with open(whitelist_save_path, 'w') as f:
                json.dump(self.final_whitelist, f, indent=4)
            logging.info(f"Saved final feature whitelist ({len(self.final_whitelist)} features) to {whitelist_save_path}")
        except Exception as e:
            logging.error(f"Failed to save final feature whitelist: {e}", exc_info=True)
            # Decide if this is critical - maybe abort if we can't save it?
            pass

        # Prune all data splits using the final whitelist
        logging.info(f"Pruning feature sets using final whitelist: {self.final_whitelist}")
        self.X_train_pruned = self.feature_engineer.prune_features(self.X_train_raw, self.final_whitelist)
        self.X_val_pruned = self.feature_engineer.prune_features(self.X_val_raw, self.final_whitelist)
        self.X_test_pruned = self.feature_engineer.prune_features(self.X_test_raw, self.final_whitelist)

        logging.info(f"Feature shapes after pruning: Train={self.X_train_pruned.shape}, Val={self.X_val_pruned.shape}, Test={self.X_test_pruned.shape}")

        # Verify all splits have the same columns after pruning
        if not (self.X_train_pruned.columns.equals(self.X_val_pruned.columns) and 
                self.X_train_pruned.columns.equals(self.X_test_pruned.columns)):
            logging.error("Column mismatch between pruned data splits. Check pruning logic.")
            # Log details for debugging
            logging.error(f"Train cols: {self.X_train_pruned.columns.tolist()}")
            logging.error(f"Val cols:   {self.X_val_pruned.columns.tolist()}")
            logging.error(f"Test cols:  {self.X_test_pruned.columns.tolist()}")
            sys.exit(1)
            
        # Check if feature sets are empty after pruning
        if self.X_train_pruned.empty or self.X_val_pruned.empty or self.X_test_pruned.empty:
             logging.error("One or more feature splits are empty after pruning. Exiting.")
             sys.exit(1)

    def scale_features(self):
        """Scales features using StandardScaler fitted on the training set."""
        logging.info("--- Stage: Scaling Features ---")
        if self.X_train_pruned is None or self.X_val_pruned is None or self.X_test_pruned is None:
            logging.error("Pruned feature sets not available for scaling.")
            sys.exit(1)

        scaler_path = os.path.join(self.current_run_models_dir, f'feature_scaler_{self.run_id}.joblib')
        
        # Ensure we only scale numeric columns
        numeric_cols = self.X_train_pruned.select_dtypes(include=np.number).columns
        if len(numeric_cols) < self.X_train_pruned.shape[1]:
            non_numeric_cols = self.X_train_pruned.select_dtypes(exclude=np.number).columns
            logging.warning(f"Non-numeric columns detected in pruned features: {non_numeric_cols.tolist()}. These will not be scaled.")
            # If non-numeric columns exist, they should ideally be handled earlier (e.g., encoding) or excluded.
            # For now, we proceed by scaling only numeric ones, but this might indicate an issue.

        if not numeric_cols.empty:
            # Check if scaler was loaded previously (when loading GRU)
            if self.scaler is None:
                logging.info("Fitting StandardScaler on training data (numeric columns only)...")
                self.scaler = StandardScaler()
                self.scaler.fit(self.X_train_pruned[numeric_cols])
                
                # Save the fitted scaler
                try:
                    joblib.dump(self.scaler, scaler_path)
                    logging.info(f"Feature scaler saved to {scaler_path}")
                except Exception as e:
                     logging.error(f"Failed to save feature scaler: {e}")
            else:
                logging.info("Using pre-loaded scaler for feature scaling.")

            # Apply scaling to all splits (numeric columns only)
            # Create copies to store scaled data, preserving original pruned dataframes
            self.X_train_scaled = self.X_train_pruned.copy()
            self.X_val_scaled = self.X_val_pruned.copy()
            self.X_test_scaled = self.X_test_pruned.copy()

            self.X_train_scaled[numeric_cols] = self.scaler.transform(self.X_train_pruned[numeric_cols])
            self.X_val_scaled[numeric_cols] = self.scaler.transform(self.X_val_pruned[numeric_cols])
            self.X_test_scaled[numeric_cols] = self.scaler.transform(self.X_test_pruned[numeric_cols])
            logging.info("Features scaled successfully.")
        else:
             logging.warning("No numeric columns found to scale. Skipping scaling step.")
             # Assign unscaled data to scaled variables to allow pipeline continuation
             self.X_train_scaled = self.X_train_pruned
             self.X_val_scaled = self.X_val_pruned
             self.X_test_scaled = self.X_test_pruned

    def run_baseline_checks(self):
        """(Optional) Runs baseline model checks."""
        logging.info("--- Stage: Baseline Checks (Placeholder) ---")
        # Placeholder - Implement if needed, e.g., LogReg on minimal features
        logging.warning("Baseline checks stage not implemented.")

    def create_sequences(self):
        """Creates sequences for GRU input using scaled features and aligned targets."""
        logging.info("--- Stage: Creating Sequences ---")
        if self.X_train_scaled is None or self.y_train is None or \
           self.X_val_scaled is None or self.y_val is None or \
           self.X_test_scaled is None or self.y_test is None:
             logging.error("Scaled features or aligned targets not available for sequence creation.")
             sys.exit(1)

        lookback = self.config['gru'].get('lookback', 60)
        horizon = self.config['gru'].get('prediction_horizon', 5) # Needed to identify target columns
        target_ret_col = f'fwd_log_ret_{horizon}'
        target_dir_col = f'direction_label_{horizon}'
        
        logging.info(f"Creating sequences with lookback={lookback}")

        # Helper function adapted from run_pipeline.py
        def _create_sequences_helper(features_scaled_df, targets_df, lookback, ret_col, dir_col):
            # Convert DataFrames to numpy arrays for efficiency
            features_np = features_scaled_df.values
            # Select target columns and convert to numpy
            y_ret_np = targets_df[ret_col].values
            y_dir_np = targets_df[dir_col].values
            
            X, y_ret_seq, y_dir_seq = [], [], []
            # Store original indices corresponding to the *target* timestep
            target_indices = [] 

            # Iterate from lookback index up to the length of the features
            for i in range(lookback, len(features_np)):
                # Append the sequence of features [i-lookback : i]
                X.append(features_np[i-lookback : i])
                # Append the target value at timestep i-1 (target corresponds to the period *ending* at i)
                # The features X[i-lookback : i] predict the target at time i.
                y_ret_seq.append(y_ret_np[i]) 
                y_dir_seq.append(y_dir_np[i])
                # Store the index of the target timestep 'i'
                target_indices.append(targets_df.index[i])

            if not X: # Check if any sequences were created
                return None, None, None, None

            # Convert lists to numpy arrays
            X_np = np.array(X)
            y_ret_seq_np = np.array(y_ret_seq)
            y_dir_seq_np = np.array(y_dir_seq)
            target_indices_pd = pd.Index(target_indices) # Keep as pandas Index
            
            return X_np, y_ret_seq_np, y_dir_seq_np, target_indices_pd

        # Create sequences for train, validation, and test sets
        self.X_train_seq, y_ret_train_seq, y_dir_train_seq, self.train_indices = _create_sequences_helper(
            self.X_train_scaled, self.y_train, lookback, target_ret_col, target_dir_col
        )
        self.X_val_seq, y_ret_val_seq, y_dir_val_seq, self.val_indices = _create_sequences_helper(
            self.X_val_scaled, self.y_val, lookback, target_ret_col, target_dir_col
        )
        self.X_test_seq, y_ret_test_seq, y_dir_test_seq, self.test_indices = _create_sequences_helper(
            self.X_test_scaled, self.y_test, lookback, target_ret_col, target_dir_col
        )

        # Check if sequences were created successfully
        if self.X_train_seq is None or self.X_val_seq is None:
             logger.error(f"Sequence creation resulted in empty train or val arrays. Check lookback ({lookback}) vs split sizes. Aborting.")
             sys.exit(1)

        logging.info(f"Sequence shapes created:")
        logging.info(f"  Train: X={self.X_train_seq.shape}, y_ret={y_ret_train_seq.shape}, y_dir={y_dir_train_seq.shape}")
        logging.info(f"  Val:   X={self.X_val_seq.shape}, y_ret={y_ret_val_seq.shape}, y_dir={y_dir_val_seq.shape}")
        logging.info(f"  Test:  X={self.X_test_seq.shape if self.X_test_seq is not None else 'None'}, y_ret={y_ret_test_seq.shape if y_ret_test_seq is not None else 'None'}, y_dir={y_dir_test_seq.shape if y_dir_test_seq is not None else 'None'}")

        # Prepare target dictionaries required by the GRU model's training and evaluation
        self.y_train_seq_dict = {
            "ret": y_ret_train_seq,
            "gauss_params": y_ret_train_seq, # Use ret target for NLL too
            "dir": y_dir_train_seq
        }
        self.y_val_seq_dict = {
            "ret": y_ret_val_seq,
            "gauss_params": y_ret_val_seq,
            "dir": y_dir_val_seq
        }
        # Test targets dictionary (useful for later evaluation/backtesting)
        if y_ret_test_seq is not None and y_dir_test_seq is not None:
             self.y_test_seq_dict = {
                 "ret": y_ret_test_seq,
                 "gauss_params": y_ret_test_seq,
                 "dir": y_dir_test_seq
             }
        else:
             self.y_test_seq_dict = None
             logging.warning("Test sequences or targets could not be created. Backtesting might fail.")

    def train_or_load_gru(self):
        """Trains a new GRU model or loads a pre-trained one using GRUModelHandler."""
        logging.info("--- Stage: Training or Loading GRU Model ---")
        gru_cfg = self.config['gru']
        train_gru_flag = self.config['control'].get('train_gru', False)

        if train_gru_flag:
            logging.info(f"Attempting to train a new GRU model for run {self.run_id}...")
            if self.X_train_seq is None or self.y_train_seq_dict is None or \
               self.X_val_seq is None or self.y_val_seq_dict is None:
                logging.error("Sequence data (train/val) not available for GRU training. Exiting.")
                sys.exit(1)

            # Get parameters from config
            lookback = gru_cfg.get('lookback', 60)
            # Get feature count from scaled data (use shape[2] for sequences)
            n_features = self.X_train_seq.shape[2]
            epochs = gru_cfg.get('epochs', 25)
            batch_size = gru_cfg.get('batch_size', 128)
            patience = gru_cfg.get('patience', 5) # Use gru patience from config

            # Train the model
            self.gru_model, history = self.gru_handler.train(
                X_train=self.X_train_seq,
                y_train_dict=self.y_train_seq_dict,
                X_val=self.X_val_seq,
                y_val_dict=self.y_val_seq_dict,
                lookback=lookback,
                n_features=n_features,
                max_epochs=epochs,
                batch_size=batch_size,
                patience=patience
            )

            if self.gru_model is None:
                logging.error("GRU model training failed. Exiting.")
                sys.exit(1)
            else:
                # Save the newly trained model
                saved_path = self.gru_handler.save() # Uses run_id from handler
                if saved_path:
                    logging.info(f"Newly trained GRU model saved to {saved_path}")
                else:
                    logging.warning("Failed to save the newly trained GRU model.")
                # Set the loaded ID to the current run ID
                self.gru_model_run_id_loaded_from = self.run_id
                logging.info(f"Using GRU model trained in current run: {self.run_id}")

        else: # Load pre-trained GRU model
            load_run_id = gru_cfg.get('model_load_run_id', None)
            if not load_run_id:
                logging.error("train_gru is False, but no gru.model_load_run_id specified in config. Exiting.")
                sys.exit(1)
            
            logging.info(f"Attempting to load pre-trained GRU model from run ID: {load_run_id}")
            # Construct the expected path using the base models directory
            model_filename = f'gru_model_{load_run_id}.keras' # Assuming .keras extension
            model_path = os.path.join(self.base_models_dir_path, f'run_{load_run_id}', model_filename)
            
            # Load the model using the handler
            self.gru_model = self.gru_handler.load(model_path)

            if self.gru_model is None:
                logging.error(f"Failed to load GRU model from path: {model_path}. Exiting.")
                sys.exit(1)
            else:
                self.gru_model_run_id_loaded_from = load_run_id
                logging.info(f"Successfully loaded GRU model from run: {load_run_id}")
                
                # --- Try loading associated scaler --- #
                scaler_filename = f'feature_scaler_{load_run_id}.joblib'
                scaler_load_path = os.path.join(self.base_models_dir_path, f'run_{load_run_id}', scaler_filename)
                logging.info(f"Attempting to load associated scaler from: {scaler_load_path}")
                if os.path.exists(scaler_load_path):
                    try:
                        self.scaler = joblib.load(scaler_load_path)
                        logging.info("Associated feature scaler loaded successfully.")
                        # --- Re-apply scaling using the loaded scaler --- #
                        numeric_cols = self.X_train_pruned.select_dtypes(include=np.number).columns
                        if not numeric_cols.empty:
                            # Important: Scale the _pruned data before sequence creation if scaler is loaded
                            logging.info("Re-scaling features using loaded scaler...")
                            self.X_train_scaled = self.X_train_pruned.copy()
                            self.X_val_scaled = self.X_val_pruned.copy()
                            self.X_test_scaled = self.X_test_pruned.copy()
                            self.X_train_scaled[numeric_cols] = self.scaler.transform(self.X_train_pruned[numeric_cols])
                            self.X_val_scaled[numeric_cols] = self.scaler.transform(self.X_val_pruned[numeric_cols])
                            self.X_test_scaled[numeric_cols] = self.scaler.transform(self.X_test_pruned[numeric_cols])
                            logging.info("Features re-scaled successfully.")
                            # Need to recreate sequences after re-scaling!
                            self.create_sequences() 
                        else:
                            logging.warning("Loaded scaler, but no numeric columns found in pruned data to re-scale.")
                    except Exception as e:
                        logging.error(f"Failed to load or apply associated scaler: {e}. Scaling might be inconsistent. Exiting.")
                        sys.exit(1)
                else:
                    logging.error(f"Associated feature scaler not found at {scaler_load_path} for run {load_run_id}. Cannot proceed. Exiting.")
                    sys.exit(1)
                # --- End Scaler Loading/Applying --- #

        # Final check: Ensure a GRU model is loaded/trained
        if self.gru_model is None:
            logging.error("No GRU model is available after train/load step. Exiting.")
            sys.exit(1)

    def calibrate_probabilities(self):
        """Calibrates GRU output probabilities using temperature scaling."""
        logging.info("--- Stage: Calibrating Probabilities ---")
        if self.gru_model is None:
            logging.error("GRU model not available for calibration. Exiting.")
            sys.exit(1)
        if self.X_val_seq is None or self.y_val_seq_dict is None:
             logging.error("Validation sequence data not available for calibration. Exiting.")
             sys.exit(1)

        # Check if a calibration temp file exists for the loaded GRU model run
        loaded_T = None
        if self.gru_model_run_id_loaded_from:
            temp_filename = f'calibration_temp_{self.gru_model_run_id_loaded_from}.npy'
            temp_load_path = os.path.join(self.base_models_dir_path, f'run_{self.gru_model_run_id_loaded_from}', temp_filename)
            if os.path.exists(temp_load_path):
                try:
                    loaded_T = np.load(temp_load_path)
                    logging.info(f"Loaded calibration temperature T={loaded_T:.4f} from GRU run {self.gru_model_run_id_loaded_from}.")
                    self.optimal_T = float(loaded_T) # Store the loaded value
                except Exception as e:
                    logging.warning(f"Failed to load calibration temp from {temp_load_path}: {e}. Recalculating.")
                    loaded_T = None # Ensure recalculation happens
            else:
                logging.info(f"No existing calibration temperature found for run {self.gru_model_run_id_loaded_from} at {temp_load_path}.")

        # If temperature wasn't loaded, calculate it using the validation set
        if loaded_T is None:
            logging.info("Calculating optimal temperature on validation set...")
            # Get predictions on validation set
            predictions_val = self.gru_handler.predict(self.X_val_seq)
            if predictions_val is None or len(predictions_val) < 3:
                 logging.error("Failed to get validation predictions for calibration. Exiting.")
                 sys.exit(1)
            
            # Extract raw probabilities P(dir=up) - assuming it's the 3rd output
            p_raw_val = predictions_val[2].flatten()
            y_dir_val = self.y_val_seq_dict['dir']

            # Check for length mismatch
            if len(p_raw_val) != len(y_dir_val):
                 logging.error(f"Mismatch between validation predictions ({len(p_raw_val)}) and targets ({len(y_dir_val)}) for calibration. Exiting.")
                 sys.exit(1)

            # Optimize temperature using the Calibrator instance
            self.optimal_T = self.calibrator.optimise_temperature(p_raw_val, y_dir_val)
            
            # Save the newly calculated temperature for the *current* run ID
            temp_save_path = os.path.join(self.current_run_models_dir, f'calibration_temp_{self.run_id}.npy')
            try:
                np.save(temp_save_path, self.optimal_T)
                logging.info(f"Saved newly calculated calibration temperature T={self.optimal_T:.4f} to {temp_save_path}")
            except Exception as e:
                logging.error(f"Failed to save calibration temperature: {e}")

        # Store the final temperature in the calibrator instance as well
        self.calibrator.optimal_T = self.optimal_T

        # Optional: Generate reliability curve plot for validation set
        if self.config.get('control', {}).get('generate_plots', True) and 'p_raw_val' in locals():
            logging.info("Generating validation reliability curve plot...")
            results_plot_dir = self.dirs.get('results', None)
            if results_plot_dir:
                 rel_curve_path = os.path.join(results_plot_dir, f'reliability_curve_val_{self.run_id}.png')
                 try:
                      # Pass calibrated probabilities using the found optimal_T
                      p_cal_val = self.calibrator.calibrate(p_raw_val) 
                      y_dir_val = self.y_val_seq_dict['dir'] # Ensure y_dir_val is available
                      # Pass save_path to the method
                      self.calibrator.reliability_curve(
                           p_pred=p_cal_val, 
                           y_true=y_dir_val,
                           plot_title=f"Reliability Curve (Validation, T={self.optimal_T:.2f})",
                           save_path=rel_curve_path
                      )
                 except Exception as e:
                      logging.error(f"Failed to generate validation reliability curve: {e}", exc_info=True)
            else:
                 logging.warning("Results directory not found, cannot save reliability curve plot.")

    def train_or_load_sac(self):
        """Trains a new SAC agent offline or loads a pre-trained one for backtesting."""
        logging.info("--- Stage: Training or Loading SAC Agent ---")
        train_sac_flag = self.config['control'].get('train_sac', False)

        if train_sac_flag:
            if self.gru_model_run_id_loaded_from is None:
                logging.error("Cannot run SAC training: GRU model run ID is not set (no model trained or loaded). Aborting.")
                sys.exit(1)
            
            logging.info(f"SAC training is enabled. Instantiating SACTrainer...")
            # Instantiate SACTrainer, passing necessary base directories from the main pipeline
            # Ensure logs/results dirs exist
            base_logs = self.dirs.get('logs')
            if not base_logs: 
                 base_logs = os.path.join(project_root, 'logs')
                 os.makedirs(base_logs, exist_ok=True)
                 logging.warning(f"Using default base logs dir for SACTrainer: {base_logs}")
                 
            base_results = self.dirs.get('results')
            if not base_results:
                 base_results = os.path.join(project_root, 'results')
                 os.makedirs(base_results, exist_ok=True)
                 logging.warning(f"Using default base results dir for SACTrainer: {base_results}")
                 
            self.sac_trainer = SACTrainer(
                config=self.config, # Pass the full config
                base_models_dir=self.base_models_dir_path,
                base_logs_dir=base_logs,
                base_results_dir=base_results 
            )
            
            # Start the training process
            final_agent_path = self.sac_trainer.train(gru_run_id_for_sac=self.gru_model_run_id_loaded_from)
            
            if final_agent_path:
                logger.info(f"SAC training completed. Final agent saved at: {final_agent_path}")
                # Set the agent path to the newly trained agent for subsequent backtesting
                self.sac_agent_load_path = final_agent_path
            else:
                logger.error("SAC training failed. Proceeding without a newly trained agent.")
                # Decide whether to fallback to loading or abort? Fallback for now.
                self.sac_agent_load_path = self._determine_sac_load_path_from_config()
                if self.sac_agent_load_path:
                     logger.warning(f"Falling back to loading SAC agent specified in config: {self.sac_agent_load_path}")
                else:
                     logger.error("SAC training failed and no load path specified in config. Cannot proceed with backtesting.")
                     # Optionally exit: sys.exit(1) 
                     # For now, allow pipeline to continue, backtester should handle None path

        else: # Load SAC agent based on config for backtesting
            logging.info("SAC training is disabled (train_sac=False). Determining agent path to load for backtesting...")
            self.sac_agent_load_path = self._determine_sac_load_path_from_config()
            if self.sac_agent_load_path:
                 logger.info(f"SAC agent path for backtesting set to load from: {self.sac_agent_load_path}")
            else:
                 logger.warning("No 'sac_load_run_id' specified in config. Backtester will need to handle using untrained/initial weights.")

    def _determine_sac_load_path_from_config(self) -> str | None:
        """Helper to determine the SAC agent load path based on config control flags."""
        load_run_id = self.config['control'].get('sac_load_run_id')
        load_step = self.config['control'].get('sac_load_step', 'final')
        sac_agent_path = None
        if load_run_id:
            # Construct path assuming structure like: <models_base>/<sac_train_run_id>/agent_<step>.pt
            # The sac_train_run_id usually differs from the pipeline run_id
            models_base = self.base_models_dir_path # Use the stored base models path
            # Assume the SAC trainer saves checkpoints inside its own run folder (e.g., models/sac_train_.../sac_agent_final)
            if load_step == 'final':
                # SAC trainer saves final model in a folder named 'sac_agent_final'
                sac_agent_path = os.path.join(models_base, load_run_id, 'sac_agent_final')
            else:
                # SAC trainer saves step checkpoints in folder 'sac_agent_step_N'
                sac_agent_path = os.path.join(models_base, load_run_id, f'sac_agent_step_{load_step}')
            
            # Check if the determined path exists
            if not os.path.exists(sac_agent_path):
                 logger.warning(f"Determined SAC load path does not exist: {sac_agent_path}. Will proceed without loading specified agent.")
                 sac_agent_path = None # Reset path if not found
                 
        return sac_agent_path

    def run_backtest(self):
        """Runs the backtest using the trained/loaded models and test data."""
        logging.info("--- Stage: Running Backtest ---")
        if not self.config['control'].get('run_backtest', False):
            logging.warning("Skipping backtest stage as per config.")
            return
        
        # Check if necessary data is available
        if self.X_test_seq is None or self.y_test_seq_dict is None or self.test_indices is None:
             logger.error("Test sequence data (X, y, indices) not available for backtesting. Skipping.")
             return
        if self.gru_handler is None or self.calibrator is None:
             logger.error("GRU Handler or Calibrator not initialized. Skipping backtest.")
             return
        if self.optimal_T is None:
             logger.warning("Optimal calibration temperature not set. Calibration might be incorrect. Proceeding...")
             # Calibrator will default to T=1.0 if its optimal_T is None
        
        # Extract original prices for plotting (need to align with test sequence indices)
        original_prices = None
        if self.y_test is not None:
            # Assuming y_test index aligns with X_test_raw index before sequencing
            # We need the price *at* the time of the target prediction
            # test_indices maps to the target time step `i` in the sequence loop
            # Need the price from the *original* test split DF aligned with these indices
            if 'close' in self.df_test.columns: # df_test has original columns before feature selection
                try:
                    original_prices = self.df_test.loc[self.test_indices, 'close']
                except KeyError:
                     logger.warning("Could not align original close prices with test indices. Price plot will be limited.")
                except Exception as e:
                    logger.error(f"Error aligning original prices: {e}")
            else:
                logger.warning("'close' column not found in df_test. Cannot extract original prices for plotting.")

        # Run the backtest using the Backtester instance
        self.backtest_results_df, self.backtest_metrics = self.backtester.run_backtest(
            sac_agent_load_path=self.sac_agent_load_path,
            X_test_seq=self.X_test_seq,
            y_test_seq_dict=self.y_test_seq_dict,
            test_indices=self.test_indices,
            gru_handler=self.gru_handler,
            calibrator=self.calibrator,
            original_prices=original_prices
        )

        if self.backtest_results_df is None or self.backtest_metrics is None:
            logger.error("Backtesting failed to produce results.")
        else:
            logger.info("Backtest completed successfully.")

    def save_results(self):
        """Saves backtest results, metrics, and plots using the Backtester instance."""
        logging.info("--- Stage: Saving Results ---")
        if self.backtest_results_df is None or self.backtest_metrics is None:
            logging.warning("No backtest results available to save. Skipping.")
            return
            
        results_dir = self.dirs.get('results')
        if not results_dir:
             logger.error("Results directory not configured. Cannot save backtest results.")
             return
             
        # Pass results to the backtester's save method
        self.backtester.save_results(
            results_df=self.backtest_results_df,
            metrics=self.backtest_metrics,
            results_dir=results_dir,
            run_id=self.run_id
        )

    def execute(self):
        """Executes the pipeline stages sequentially."""
        logging.info("=== Starting Pipeline Execution ===")
        try:
            self.load_and_preprocess_data()
            self.engineer_features()
            self.define_labels_and_align()
            self.split_data()
            self.select_and_prune_features()
            # Scaling is now handled conditionally within train_or_load_gru if loading
            # If training, scaling happens afterwards based on fitted scaler.
            # self.scale_features() # Moved
            self.train_or_load_gru() # Handles loading/training GRU and associated scaler/rescaling
            # If GRU was trained, scale features *now* using the fitted scaler
            if self.config['control'].get('train_gru', False):
                 self.scale_features()
                 # Need to recreate sequences *after* scaling if GRU was trained
                 self.create_sequences()
            self.run_baseline_checks() # Optional
            # self.create_sequences() # Sequence creation now happens conditionally after scaling
            self.calibrate_probabilities()
            self.train_or_load_sac()
            self.run_backtest() # Call the implemented backtest method
            self.save_results() # Call the implemented save method
            logging.info("=== Pipeline Execution Finished Successfully ===")

        except Exception as e:
            logging.error(f"Pipeline execution failed: {e}", exc_info=True)
            logging.error("=== Pipeline Execution Terminated Due to Error ===")
            sys.exit(1)

# --- Entry Point --- #

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run the GRU-SAC Trading Pipeline.")
    # Default config path relative to the project root/gru_sac_predictor level
    # Adjusting default config path seeking strategy
    # 1. Relative to project root (../config.yaml from src)
    default_config_rel_root = os.path.abspath(os.path.join(script_dir, '..', 'config.yaml'))
    # 2. Relative to package dir (../gru_sac_predictor/config.yaml from project root)
    default_config_pkg = os.path.join(os.path.dirname(project_root), 'gru_sac_predictor', 'config.yaml')
    # 3. Relative to CWD
    default_config_cwd = os.path.abspath('config.yaml')

    # Determine the best default path
    if os.path.exists(default_config_rel_root):
        default_config = default_config_rel_root
    elif os.path.exists(default_config_pkg):
        default_config = default_config_pkg
    else:
        default_config = default_config_cwd # Fallback to CWD
        if not os.path.exists(default_config):
             # If none exist, use a placeholder path to show in help message
             default_config = '../config.yaml' # Placeholder for help message

    parser.add_argument(
        '--config',
        type=str,
        default=default_config,
        help=f"Path to the configuration YAML file (default attempts: {default_config_rel_root}, {default_config_pkg}, {default_config_cwd})"
    )
    args = parser.parse_args()

    config_to_use = args.config
    # Verify the provided/default config path exists
    if not os.path.exists(config_to_use):
        print(f"Error: Config file not found at the specified path: {config_to_use}")
        print("Please ensure the path is correct or place config.yaml in the expected location.")
        sys.exit(1)

    # Instantiate and run the pipeline
    pipeline = TradingPipeline(config_path=config_to_use)
    pipeline.execute() 