"""
Handles GRU Model Training, Loading, Saving, and Prediction.

This is the comprehensive GRU model implementation that includes both v2 and v3
model architectures as well as the handler for training, saving, loading, and prediction.
"""

import sys
import os
import json
import joblib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import logging
from tqdm.keras import TqdmCallback
from typing import Dict, Tuple, Any, List, Optional, Callable

# --- Add Imports for TF/Keras and Optional TFA --- #
try:
    import tensorflow as tf
    from tensorflow.keras import Model, layers, callbacks, saving
    from tensorflow.keras.losses import Huber, CategoricalCrossentropy
    from tensorflow.keras.optimizers import Adam
    from tensorflow.keras import regularizers
    from tensorflow.keras.metrics import MeanAbsoluteError, RootMeanSquaredError # type: ignore
    KERAS_AVAILABLE = True
except ImportError:
    logger.error("TensorFlow/Keras not found. GRU model functionality will be unavailable.")
    # Define placeholders for types if Keras is not available
    Model = Any
    layers = Any
    callbacks = Any
    saving = Any
    Huber = Any
    Adam = Any
    regularizers = Any
    MeanAbsoluteError = Any
    RootMeanSquaredError = Any
    KERAS_AVAILABLE = False

# --- Removed TFA Import Logic --- #
# try:
#     import tensorflow_addons as tfa
#     # Check for specific loss needed (optional, but good practice)
#     if hasattr(tfa, 'losses') and (hasattr(tfa.losses, 'BinaryFocalCrossentropy') or hasattr(tfa.losses, 'CategoricalFocalCrossentropy')):
#         TFA_AVAILABLE = True
#         logger.info("TensorFlow Addons found and contains required Focal Loss classes.")
#     else:
#         TFA_AVAILABLE = False
#         logger.warning("TensorFlow Addons found, but required Focal Loss classes (Binary/Categorical) seem missing.")
# except ImportError:
#     logger.warning("TensorFlow Addons (tfa) not found. Focal Loss will not be available.")
#     tfa = None # Define tfa as None if not imported
#     TFA_AVAILABLE = False
# --- End Removed TFA Import Logic --- #

# ===================================================================
#  UTILITIES AND LOSS FUNCTIONS
# ===================================================================

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

# --- Define create_mask function globally --- #
def create_mask(inputs):
    """Creates a lower-triangular causal mask for Attention layers."""
    input_shape = tf.shape(inputs)
    # Need batch_size only if modifying mask per batch, otherwise seq_len is enough
    # batch_size = input_shape[0]
    seq_len = input_shape[1]
    # Lower triangular mask (ones below and on diagonal, zeros above)
    mask = tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)
    # MHA expects mask shape like (batch_size, seq_len, seq_len) or (seq_len, seq_len)
    # Returning (seq_len, seq_len) is usually sufficient as Keras broadcasts.
    return mask
# --- End global function definition ---

# --- Custom Focal Loss Implementations --- #

# --- Removed Binary Focal Loss --- #
# def binary_focal_loss(y_true, y_pred, gamma=2.0, alpha=0.25):
#     ...
# --- End Removed Binary Focal Loss --- #

def categorical_focal_loss(y_true, y_pred, gamma=2.0, alpha=0.25, label_smoothing=0.0):
    """Categorical focal loss implementation.
    Handles one-hot encoded targets and optional label smoothing.
    Source: Adapted from https://www.tensorflow.org/api_docs/python/tf/keras/losses/CategoricalFocalCrossentropy
    """
    if not KERAS_AVAILABLE:
        raise ImportError("Keras/TensorFlow backend needed for focal loss calculation.")

    y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.cast(y_true, y_pred.dtype)

    # Apply label smoothing
    if label_smoothing > 0.0:
        num_classes = tf.cast(tf.shape(y_true)[-1], y_pred.dtype)
        y_true = y_true * (1.0 - label_smoothing) + (label_smoothing / num_classes)

    # Calculate cross-entropy loss
    epsilon = tf.keras.backend.epsilon()
    y_pred = tf.clip_by_value(y_pred, epsilon, 1.0 - epsilon)
    ce = -y_true * tf.math.log(y_pred)
    
    # Calculate focal loss components
    p_t = tf.reduce_sum(y_true * y_pred, axis=-1) # Probability of the true class
    focal_factor = tf.pow(1.0 - p_t, gamma)
    
    # Weighted loss
    # Note: TF/Keras implementation uses alpha on per-class basis if provided as array.
    # Here, we use a single alpha for simplicity, but could be adapted.
    # If alpha is used, it typically applies to the positive class contribution.
    # For categorical, it's less common or needs careful definition. Let's omit alpha here.
    # weighted_loss = alpha * focal_factor * ce
    focal_loss_per_example = focal_factor * tf.reduce_sum(ce, axis=-1)

    return tf.reduce_mean(focal_loss_per_example)
# --- End Custom Focal Loss --- #

@saving.register_keras_serializable(package='GRU')
def gaussian_nll(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
    """Gaussian negative‑log likelihood for *scalar* targets.

    The model is assumed to predict concatenated [mu, log_sigma].

    Given targets :math:`y` and predictions :math:`\mu, \log\sigma`,
    the NLL is

    .. math::

        \mathcal{L}_{NLL} = \frac{1}{2} \exp(-2\log\sigma)(y-\mu)^2 
                             + \log\sigma.

    A small constant is added for numerical stability.
    """
    mu, log_sigma = tf.split(y_pred, 2, axis=-1)
    # Ensure y_true has the same shape as mu for the subtraction
    y_true_shaped = tf.reshape(y_true, tf.shape(mu))
    inv_var = tf.exp(-2.0 * log_sigma)  # = 1/sigma^2
    nll = 0.5 * inv_var * tf.square(y_true_shaped - mu) + log_sigma
    return tf.reduce_mean(nll)

# ===================================================================
#  MODEL BUILDERS
# ===================================================================

# --- Removed build_gru_model (v2) --- #
# def build_gru_model(lookback: int, n_features: int, kappa: float = 0.2) -> Model:
#     ...
# --- End Removed build_gru_model (v2) --- #

def build_gru_model_v3(
    lookback: int,
    n_features: int,
    gru_units: int = 96,
    attention_units: int = 16,
    dropout_rate: float = 0.1,
    learning_rate: float = 1e-4,
    focal_gamma: float = 2.0,
    focal_label_smoothing: float = 0.1,
    huber_delta: float = 1.0,
    loss_weight_mu: float = 0.3,
    loss_weight_dir3: float = 1.0,
    l2_reg: float = 1e-4
) -> Model:
    """
    Builds and compiles the GRU v3 model based on the specified architecture and hyperparameters.

    Architecture: Input -> GRU -> LayerNorm -> Attention -> LayerNorm -> Output Heads

    Args:
        lookback (int): The sequence length for the GRU input.
        n_features (int): The number of features at each timestep.
        gru_units (int): Number of units for the GRU layer.
        attention_units (int): Number of units for the Attention layer.
        dropout_rate (float): Dropout rate for the GRU and Attention layers.
        learning_rate (float): Learning rate for the Adam optimizer.
        focal_gamma (float): Gamma parameter for CategoricalFocalCrossentropy.
        focal_label_smoothing (float): Label smoothing for CategoricalFocalCrossentropy.
        huber_delta (float): Delta parameter for Huber loss.
        loss_weight_mu (float): Weight for the 'mu' output loss.
        loss_weight_dir3 (float): Weight for the 'dir3' output loss.
        l2_reg (float): L2 regularization factor for dense layers.

    Returns:
        keras.Model: The compiled Keras model.
    """
    
    # --- Ensure l2_reg is a float --- #
    try:
        l2_reg_float = float(l2_reg)
    except (ValueError, TypeError) as e:
        logger.error(f"Invalid value for l2_reg: '{l2_reg}'. Expected a number. Error: {e}. Using 0.0.")
        l2_reg_float = 0.0
    # --- End Ensure --- #

    # --- Ensure learning_rate is a float --- #
    try:
        learning_rate_float = float(learning_rate)
    except (ValueError, TypeError) as e:
        logger.error(f"Invalid value for learning_rate: '{learning_rate}'. Expected a number. Error: {e}. Using 1e-4.")
        learning_rate_float = 1e-4 # Default value
    # --- End Ensure --- #

    input_shape = (lookback, n_features)
    inputs = layers.Input(shape=input_shape)
    
    # GRU Layer + Layer Norm (Revision 2-C)
    x = layers.GRU(
        gru_units, 
        return_sequences=True, 
        name='gru_base', 
        reset_after=True  # Force standard kernel implementation
    )(inputs)
    x = layers.LayerNormalization(name='gru_layernorm')(x) # Added LayerNorm
    x = layers.Dropout(dropout_rate, name='dropout_gru')(x)
    
    # Attention Layer (if attention_units > 0)
    if attention_units > 0:
        # --- Create Causal Mask using Lambda Layer (Keras Compatible) --- # 
        causal_mask = layers.Lambda(create_mask, name='causal_mask_lambda')(x)
        # --- End Causal Mask Creation --- #

        attn_output = layers.MultiHeadAttention(
            num_heads=max(1, attention_units // 16), # Example heuristic for num_heads
            key_dim=attention_units, 
            kernel_regularizer=regularizers.l2(l2_reg_float), # Use float value
            dropout=dropout_rate,
            # --- REMOVED use_causal_mask --- #
            # use_causal_mask=True, 
            # --- End REMOVED --- #
            name='multi_head_attention'
        # --- RE-ADD attention_mask argument --- #
        )(query=x, value=x, key=x, attention_mask=causal_mask) # Pass the mask tensor from Lambda
        # --- End RE-ADD --- #
        x = layers.Dropout(dropout_rate, name='dropout_attn')(attn_output)
    # If attention_units is 0, skip attention layer
    
    # Pooling Layer (Consider GlobalMaxPooling1D as alternative?)
    pooled_output = layers.GlobalAveragePooling1D(name='global_avg_pool')(x)
    
    # Dense Heads with L2 Regularization (Revision 2-C)
    kernel_regularizer = regularizers.l2(l2_reg_float) # Use float value
    
    # For the classification head, we'll separate logits and activation for potential logit extraction
    logits_dir3 = layers.Dense(
        3, 
        name='dir3_logits', 
        kernel_regularizer=kernel_regularizer
    )(pooled_output)
    dir3_output = layers.Activation('softmax', name='dir3')(logits_dir3)
    
    mu_output = layers.Dense(
        1, 
        activation='linear', 
        name='mu',
        kernel_regularizer=kernel_regularizer
    )(pooled_output)
    
    model = tf.keras.Model(inputs=inputs, outputs=[mu_output, dir3_output])

    # Compile using passed hyperparameters
    # Use actual FocalLoss if TFA is available, otherwise the standard loss
    # dir3_loss_instance = FocalLoss(gamma=focal_gamma, label_smoothing=focal_label_smoothing) if TFA_AVAILABLE else FocalLoss(label_smoothing=focal_label_smoothing)
    # --- Use local categorical_focal_loss --- #
    dir3_loss_instance = lambda y_true, y_pred: categorical_focal_loss(
        y_true, y_pred, gamma=focal_gamma, label_smoothing=focal_label_smoothing
    )
    dir3_loss_instance.__name__ = 'categorical_focal_loss' # For Keras saving/loading
    # --- End Use local --- #
    losses = {
        "dir3": dir3_loss_instance, 
        "mu": Huber(delta=huber_delta)
    }
    loss_weights = {"dir3": loss_weight_dir3, "mu": loss_weight_mu}
    optimizer = Adam(learning_rate=learning_rate_float, clipnorm=1.0)
    metrics = {"dir3": ['accuracy']} 

    try:
        model.compile(
            optimizer=optimizer,
            loss=losses,
            loss_weights=loss_weights,
            metrics=metrics
        )
        logger.info("GRU v3 model built and compiled successfully.")
    except Exception as e:
         logger.error(f"Failed to compile GRU v3 model: {e}")
         # Re-raise to prevent using uncompiled model
         raise e 

    return model

class GRUModelHandler:
    """Manages the lifecycle of the GRU model."""

    def __init__(self, run_id: str, models_dir: str, config: dict):
        """
        Initialize the handler.

        Args:
            run_id (str): The current pipeline run ID.
            models_dir (str): The base directory where models for this run are saved.
            config (dict): The pipeline configuration dictionary.
        """
        self.run_id = run_id
        self.models_dir = models_dir # Should be the specific directory for this run
        self.config = config # Store config
        self.model: Model | None = None
        self.model_version_used = None # Track which version was built/loaded
        self.use_ternary = self.config.get('gru', {}).get('use_ternary', False) # Store ternary flag
        logger.info(f"GRUModelHandler initialized for run {run_id} in {models_dir}")

    def train(
        self,
        X_train: np.ndarray,
        y_train_dict: Dict[str, np.ndarray],
        X_val: np.ndarray,
        y_val_dict: Dict[str, np.ndarray],
        lookback: int,
        n_features: int,
        max_epochs: int = 25,
        batch_size: int = 128,
        patience: int = 3
    ) -> Tuple[Model | None, Any]: # Returns model and history
        """Trains the GRU model using provided data and configuration."""
        if not KERAS_AVAILABLE:
            logger.error("Keras is not available. Cannot train GRU model.")
            return None, None

        if self.model is None:
            logger.info("No existing model found. Building a new one for training.")
            self.model = build_gru_model_v3(
                lookback=lookback,
                n_features=n_features,
                # Pass hyperparameters from config
                gru_units=self.config.get('gru', {}).get('gru_units', 96),
                attention_units=self.config.get('gru', {}).get('attention_units', 16),
                dropout_rate=self.config.get('gru', {}).get('dropout_rate', 0.1),
                learning_rate=self.config.get('gru', {}).get('learning_rate', 1e-4),
                focal_gamma=self.config.get('gru', {}).get('focal_gamma', 2.0),
                focal_label_smoothing=self.config.get('gru', {}).get('focal_label_smoothing', 0.1),
                huber_delta=self.config.get('gru', {}).get('huber_delta', 1.0),
                loss_weight_mu=self.config.get('gru', {}).get('loss_weight_mu', 0.3),
                loss_weight_dir3=self.config.get('gru', {}).get('loss_weight_dir3', 1.0),
                l2_reg=self.config.get('gru', {}).get('l2_reg', 1e-4)
            )
        else:
            logger.info("Using the existing model for training.")

        if self.model is None:
             logger.error("Failed to build or retrieve model for training.")
             return None, None

        # --- Callbacks --- #
        # 1. Early Stopping
        early_stopping = callbacks.EarlyStopping(
            monitor='val_loss', 
            patience=patience, 
            restore_best_weights=True,
            verbose=1
        )
        
        # 2. Model Checkpoint (optional, EarlyStopping restores best weights)
        # Consider saving best model directly after training finishes if needed
        # checkpoint_path = os.path.join(self.models_dir, f'best_gru_model_{self.run_id}.keras') 
        # model_checkpoint = callbacks.ModelCheckpoint(
        #     filepath=checkpoint_path,
        #     monitor='val_loss',
        #     save_best_only=True,
        #     save_weights_only=False,
        #     verbose=1
        # )
        
        # 3. TQDM Progress Bar
        tqdm_callback = TqdmCallback(verbose=1)

        # 4. CSV Logger (V3 Output Contract)
        # Define logs directory based on models_dir structure (assuming logs are sibling)
        base_dir = os.path.dirname(self.models_dir) # e.g., /path/to/models
        logs_base_dir = os.path.join(os.path.dirname(base_dir), 'logs') # /path/to/logs
        run_logs_dir = os.path.join(logs_base_dir, self.run_id) # /path/to/logs/<run_id>
        os.makedirs(run_logs_dir, exist_ok=True)
        csv_log_path = os.path.join(run_logs_dir, 'gru_history.csv')
        logger.info(f"Setting up CSVLogger to save history to: {csv_log_path}")
        csv_logger = callbacks.CSVLogger(csv_log_path, append=False)

        # --- Prepare target data structure for model.fit --- #
        # The keys in the dictionary passed to y must match the names of the output layers.
        # Model output layer names are 'mu' and 'dir3' (as defined in build_gru_model_v3)

        try:
            # Find the actual keys in the input dictionary
            input_keys = list(y_train_dict.keys())
            ret_key = next((k for k in input_keys if 'fwd_log_ret' in k), None)
            dir_key = next((k for k in input_keys if 'direction_label' in k), None)
            
            if ret_key is None or dir_key is None:
                raise KeyError(f"Could not find required target keys ('fwd_log_ret...', 'direction_label...') in y_train_dict. Found: {input_keys}")

            # Map model output names ('mu', 'dir3') to data using actual input keys
            y_train_fit = {
                'mu': y_train_dict[ret_key],
                'dir3': y_train_dict[dir_key]
            }
            # Assume y_val_dict has the same keys
            y_val_fit = {
                'mu': y_val_dict[ret_key],
                'dir3': y_val_dict[dir_key]
            }
            logger.debug(f"Target structure for model.fit prepared. Model output names mapped: {list(y_train_fit.keys())}")
        except KeyError as e:
            logger.error(f"KeyError preparing target dictionary for model.fit: {e}")
            logger.error(f"Available train keys: {list(y_train_dict.keys())}")
            logger.error(f"Available val keys: {list(y_val_dict.keys())}")
            return None, None # Indicate failure
        # --- End target preparation --- #

        # --- Handle potential NaNs/Infs before training ---
        if np.isnan(X_train).any() or np.isinf(X_train).any():
            logger.error("NaN or Inf found in X_train sequence data before training. Aborting.")
            return None, None
        if any(np.isnan(arr).any() or np.isinf(arr).any() for arr in y_train_fit.values()):
            logger.error("NaN or Inf found in y_train_fit sequence data before training. Aborting.")
            return None, None
        if np.isnan(X_val).any() or np.isinf(X_val).any():
            logger.error("NaN or Inf found in X_val sequence data before training. Aborting.")
            return None, None
        if any(np.isnan(arr).any() or np.isinf(arr).any() for arr in y_val_fit.values()):
            logger.error("NaN or Inf found in y_val_fit sequence data before training. Aborting.")
            return None, None
        logger.info("Data checks passed: No NaNs or Infs found in training/validation sequences.")
        # --- End Check ---

        logger.info(f"Starting GRU model training for {max_epochs} epochs...")
        history = self.model.fit(
            X_train, 
            y_train_fit, # Pass the prepared dictionary
            epochs=max_epochs, 
            batch_size=batch_size, 
            validation_data=(X_val, y_val_fit), # Pass the prepared dictionary 
            callbacks=[
                early_stopping, 
                # model_checkpoint, # Optional: Rely on restore_best_weights
                tqdm_callback,
                csv_logger # Add CSV Logger
            ],
            shuffle=False, # Important for time series
            verbose=1 # Or 2 for more detail per epoch
        )

        return self.model, history

    def save(self, model_name: str = 'gru_model') -> str | None:
        """
        Saves the current model to the run's model directory.
        Appends model version to the filename.
        """
        if self.model is None:
            logger.error("No model available to save.")
            return None
        if self.model_version_used is None:
            logger.warning("Model version was not set during training/loading. Saving with default name.")
            version_suffix = "unknown"
        else:
             version_suffix = self.model_version_used # e.g., 'v2' or 'v3'

        # Use .keras format and include version in filename
        save_filename = f"{model_name}_{self.run_id}.keras"
        save_path = os.path.join(self.models_dir, save_filename)
        try:
            self.model.save(save_path)
            logger.info(f"GRU model saved successfully to: {save_path}")
            return save_path
        except Exception as e:
            logger.error(f"Failed to save GRU model to {save_path}: {e}", exc_info=True)
            return None

    def load(self, model_path: str) -> Model | None:
        """
        Loads a GRU model from the specified path.
        Handles custom objects if needed (primarily for v2 gaussian_nll).
        """
        if not os.path.exists(model_path):
            logger.error(f"Model file not found at: {model_path}")
            return None

        logger.info(f"Loading GRU model from: {model_path}")
        try:
            # Custom objects needed mainly for v2's gaussian_nll
            # custom_objects = {'gaussian_nll': gaussian_nll}
            # v3 doesn't need gaussian_nll, but might need focal loss if not standard
            custom_objects = {
                'categorical_focal_loss': categorical_focal_loss,
                'FocalLoss': categorical_focal_loss, # Add alias if model was saved with class name
                # --- Add create_mask to custom objects --- #
                'create_mask': create_mask
                # --- End Add --- #
            }
            
            # Attempt to load FocalLoss if needed (REMOVED - Using local implementation)
            # if TFA_AVAILABLE:
            #     ...

            self.model = tf.keras.models.load_model(model_path, custom_objects=custom_objects)
            logger.info("GRU model loaded successfully.")
            # Try to infer model version from loaded model output names
            try:
                if 'dir3' in self.model.output_names:
                    self.model_version_used = 'v3'
                # elif 'dir' in self.model.output_names: # Assuming v2 used 'dir'
                #      self.model_version_used = 'v2' # Removed v2 check
                else:
                     self.model_version_used = 'unknown'
                logger.info(f"Inferred loaded model version: {self.model_version_used}")
            except Exception:
                 logger.warning("Could not infer model version from loaded model.")
                 self.model_version_used = 'unknown'
                 
            self.model.summary(print_fn=logger.info) # Log summary of loaded model
            return self.model
        except Exception as e:
            logger.error(f"Failed to load GRU model from {model_path}: {e}", exc_info=True)
            self.model = None
            return None

    # --- Add Logits Prediction Method (Task 4) --- #
    def _make_logits_model(self) -> Model | None:
        """Builds a frozen view that outputs dir3 pre-softmax logits."""
        if self.model is None:
            logger.error("Cannot create logits view: Main model is not loaded/built.")
            return None
            
        # Check if the view is already cached
        if hasattr(self, "_logit_view") and self._logit_view is not None:
            return self._logit_view
            
        try:
            # Check if the expected logits layer exists
            logits_layer = self.model.get_layer("dir3_logits")
            logits_tensor = logits_layer.output
            # Create a new model sharing the inputs and weights
            self._logit_view = tf.keras.Model(
                inputs=self.model.input, 
                outputs=logits_tensor,
                name="gru_logits_view" # Optional name
            )
            # No compilation needed for inference-only model
            logger.info("Created inference-only model view for 'dir3_logits' output.")
            return self._logit_view
        except ValueError:
             logger.error("Layer 'dir3_logits' not found in the current model. Cannot create logits view.")
             self._logit_view = None # Ensure cache is None
             return None
        except Exception as e:
             logger.error(f"Error creating logits view model: {e}", exc_info=True)
             self._logit_view = None
             return None

    def predict_logits(self, X_data: np.ndarray, batch_size: int = 1024) -> np.ndarray | None:
        """Returns raw logits (n,3) for Vector Scaling calibration using a model view."""
        logit_model = self._make_logits_model()
        
        if logit_model is None:
            logger.error("Logits view model is not available. Cannot predict logits.")
            return None
            
        if X_data is None or len(X_data) == 0:
             logger.warning("Input data for logit prediction is None or empty.")
             return None

        logger.info(f"Generating logit predictions for {len(X_data)} samples...")
        try:
            # Use verbose=0 to avoid Keras progress bars for this internal prediction
            logits = logit_model.predict(X_data, batch_size=batch_size, verbose=0)
            logger.info("Logit predictions generated successfully.")
            logger.debug(f"Logits output shape: {logits.shape}")
            return logits
        except Exception as e:
            logger.error(f"Error during logit prediction: {e}", exc_info=True)
            return None
    # --- End Logits Prediction Method --- #

    def predict(self, X_data: np.ndarray, batch_size: int = 1024) -> Any:
        """
        Generates predictions using the loaded/trained model.
        """
        if self.model is None:
            logger.error("No model available for prediction.")
            return None
        if X_data is None or len(X_data) == 0:
             logger.warning("Input data for prediction is None or empty.")
             return None

        logger.info(f"Generating predictions for {len(X_data)} samples...")
        try:
            predictions = self.model.predict(X_data, batch_size=batch_size)
            logger.info("Predictions generated successfully.")
            if isinstance(predictions, list):
                 pred_shapes = [p.shape for p in predictions]
                 logger.debug(f"Prediction output shapes: {pred_shapes}")
            else:
                 logger.debug(f"Prediction output shape: {predictions.shape}")
            return predictions
        except Exception as e:
            logger.error(f"Error during model prediction: {e}", exc_info=True)
            return None 

    # --- New Method for Temporary Training (Nested CV) --- #
    def build_and_train_temporary(
        self,
        gru_config: Dict[str, Any],
        X_train: np.ndarray,
        y_train_dict: Dict[str, np.ndarray],
        X_val: np.ndarray,
        y_val_dict: Dict[str, np.ndarray],
        lookback: int,
        n_features: int,
    ) -> Tuple[Optional[Model], Optional[Any]]:
        """
        Builds and trains a temporary GRU model for hyperparameter tuning trials.
        Does NOT save persistent checkpoints or modify the main handler state (self.model).

        Args:
            gru_config (dict): Hyperparameters for this specific trial.
            X_train, y_train_dict: Training sequences and targets for the inner fold.
            X_val, y_val_dict: Validation sequences and targets for the inner fold.
            lookback (int): Input sequence length.
            n_features (int): Number of input features.

        Returns:
            Tuple[Model | None, History | None]: The trained (but not saved) Keras model 
                                                 and the History object, or (None, None) on failure.
        """
        logger.info(f"--- Starting TEMPORARY GRU training (Trial) ---")
        logger.debug(f"Trial Params: {gru_config}")

        # Determine parameters from trial config
        # Use .get with defaults from the main config stored in self.config if available
        _max_epochs = int(gru_config.get('epochs', self.config.get('gru', {}).get('epochs', 25))) 
        _batch_size = int(gru_config.get('batch_size', self.config.get('gru', {}).get('batch_size', 128)))
        _patience = int(gru_config.get('patience', self.config.get('gru', {}).get('patience', 5))) 

        temp_model: Optional[Model] = None
        temp_history: Optional[Any] = None

        # 1. Build the temporary model
        try:
            # Ensure parameters are passed as correct types
            temp_model = build_gru_model_v3(
                lookback=lookback,
                n_features=n_features,
                gru_units=int(gru_config.get('gru_units', 96)),
                attention_units=int(gru_config.get('attention_units', 16)),
                dropout_rate=float(gru_config.get('dropout_rate', 0.1)),
                learning_rate=float(gru_config.get('learning_rate', 1e-4)),
                focal_gamma=float(gru_config.get('focal_gamma', 2.0)),
                focal_label_smoothing=float(gru_config.get('focal_label_smoothing', 0.1)),
                huber_delta=float(gru_config.get('huber_delta', 1.0)),
                loss_weight_mu=float(gru_config.get('loss_weight_mu', 0.3)),
                loss_weight_dir3=float(gru_config.get('loss_weight_dir3', 1.0)),
                l2_reg=float(gru_config.get('l2_reg', 1e-4))
            )
            logger.info("Temporary GRU model built.")
        except Exception as build_err:
             logger.error(f"Error building temporary GRU model: {build_err}", exc_info=True)
             return None, None

        # 2. Setup TEMPORARY Callbacks (No ModelCheckpoint or TensorBoard)
        early_stopping_callback = callbacks.EarlyStopping(
            monitor='val_loss', 
            patience=_patience, 
            verbose=0, # Less verbose for trials
            restore_best_weights=True # Important for getting best trial performance
        )
        tqdm_callback = TqdmCallback(verbose=0) # Silent progress bar

        temp_callbacks = [early_stopping_callback, tqdm_callback]

        # 3. Train the temporary model
        try:
            logger.debug(f"Starting temporary model.fit (Epochs={_max_epochs}, Batch={_batch_size})...")
            temp_history = temp_model.fit(
                X_train, 
                y_train_dict, 
                epochs=_max_epochs, 
                batch_size=_batch_size,
                validation_data=(X_val, y_val_dict), 
                callbacks=temp_callbacks,
                shuffle=False,
                verbose=0 
            )
            logger.debug("Temporary model training finished.")
            # Log best val_loss achieved in this temporary training
            if temp_history and 'val_loss' in temp_history.history:
                if temp_history.history['val_loss']: # Check if list is not empty
                    best_val_loss = min(temp_history.history['val_loss'])
                    logger.debug(f"Temporary training best val_loss: {best_val_loss:.5f}")
                else:
                    logger.warning("Temporary training history 'val_loss' list is empty.")
            else:
                logger.warning("Temporary training history missing 'val_loss' key or history is None.")
                
        except Exception as train_err:
            logger.warning(f"Error during temporary GRU model training: {train_err}") # Log as warning for trials
            temp_model = None
            temp_history = None

        return temp_model, temp_history
    # --- End New Method --- #

    def save(self, model_name: str = 'gru_model') -> str | None:
        """
        Saves the trained model architecture, weights, scaler, and metadata.
        Uses Keras standard saving format (.keras).
        """
        if self.model is None:
            logger.error("No model available to save.")
            return None
        if self.model_version_used is None:
            logger.warning("Model version was not set during training/loading. Saving with default name.")
            version_suffix = "unknown"
        else:
             version_suffix = self.model_version_used # e.g., 'v2' or 'v3'

        # Use .keras format and include version in filename
        save_filename = f"{model_name}_{self.run_id}.keras"
        save_path = os.path.join(self.models_dir, save_filename)
        try:
            self.model.save(save_path)
            logger.info(f"GRU model saved successfully to: {save_path}")
            return save_path
        except Exception as e:
            logger.error(f"Failed to save GRU model to {save_path}: {e}", exc_info=True)
            return None 