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

import tensorflow as tf
from tensorflow.keras import Model, callbacks
from tensorflow.keras import saving # Use saving module for custom object registration
import numpy as np
import os
import logging
from tqdm.keras import TqdmCallback
from typing import Dict, Tuple, Any

# Import necessary components from model_gru
# It's better if build_gru_model and gaussian_nll are defined here 
# or imported reliably. Assuming they stay in model_gru for now.
try:
    from .model_gru import build_gru_model, gaussian_nll 
except ImportError:
     # If run directly, might need adjustment or place build_gru_model here
     logging.error("Failed to import build_gru_model/gaussian_nll from .model_gru. Ensure structure is correct.")
     # Define gaussian_nll here as a fallback if import fails and it's critical
     @saving.register_keras_serializable(package='GRU')
     def gaussian_nll(y_true: tf.Tensor, y_pred: tf.Tensor) -> tf.Tensor:
         mu, log_sigma = tf.split(y_pred, 2, axis=-1)
         y_true_shaped = tf.reshape(y_true, tf.shape(mu))
         inv_var = tf.exp(-2.0 * log_sigma)
         nll = 0.5 * inv_var * tf.square(y_true_shaped - mu) + log_sigma
         return tf.reduce_mean(nll)
     # build_gru_model would also need to be defined here as a fallback
     # This indicates a potential structure issue if the import fails

logger = logging.getLogger(__name__)

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

    def __init__(self, run_id: str, models_dir: str):
        """
        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.
        """
        self.run_id = run_id
        self.models_dir = models_dir # Should be the specific directory for this run
        self.model: Model | None = None
        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
        """
        Builds and trains the GRU model.

        Args:
            X_train: Training feature sequences.
            y_train_dict: Dictionary of training targets.
            X_val: Validation feature sequences.
            y_val_dict: Dictionary of validation targets.
            lookback: Sequence length.
            n_features: Number of features per timestep.
            max_epochs: Maximum training epochs.
            batch_size: Training batch size.
            patience: Early stopping patience.

        Returns:
            Tuple[Model | None, Any]: The trained Keras model (or None on failure) and the training history object.
        """
        logger.info(f"Building GRU model: lookback={lookback}, n_features={n_features}")
        try:
            # Ensure build_gru_model is available
            if 'build_gru_model' not in globals() and 'build_gru_model' not in locals():
                 raise NameError("build_gru_model function is not defined or imported.")
            self.model = build_gru_model(lookback, n_features)
            logger.info("Model built successfully.")
            self.model.summary(print_fn=logger.info) # Log model summary
        except Exception as e:
            logger.error(f"Failed to build GRU model: {e}", exc_info=True)
            return None, None

        cb_early = callbacks.EarlyStopping(
            monitor="val_loss", # Monitor overall validation loss
            patience=patience,
            mode='min',
            restore_best_weights=True,
            verbose=1,
        )
        cb_tqdm = TqdmCallback(verbose=1)

        logger.info(f"Starting GRU training: epochs={max_epochs}, batch={batch_size}, patience={patience}")
        logger.info(f"  Train X shape: {X_train.shape}")
        logger.info(f"  Val X shape:   {X_val.shape}")
        logger.info(f"  Train y keys: {list(y_train_dict.keys())}")
        logger.info(f"  Val y keys:   {list(y_val_dict.keys())}")

        history = None
        try:
            history = self.model.fit(
                X_train,
                y_train_dict,
                validation_data=(X_val, y_val_dict),
                epochs=max_epochs,
                batch_size=batch_size,
                callbacks=[cb_early, cb_tqdm],
                verbose=0, # Let tqdm handle progress
            )
            logger.info("GRU training finished.")
            logger.info(f"Best validation loss: {min(history.history['val_loss']):.4f}")
            return self.model, history
        except Exception as e:
            logger.error(f"Error during GRU model training: {e}", exc_info=True)
            return None, history # Return None for model but maybe history is useful

    def save(self, model_name: str = 'gru_model') -> str | None:
        """
        Saves the current model to the run's model directory.

        Args:
            model_name (str): The base name for the saved model file (e.g., 'gru_model').
                               The run_id will be appended.

        Returns:
             str | None: The full path to the saved model file, or None if saving failed.
        """
        if self.model is None:
            logger.error("No model available to save.")
            return None

        # Use .keras format for modern saving
        save_path = os.path.join(self.models_dir, f"{model_name}_{self.run_id}.keras")
        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 the custom gaussian_nll loss function.

        Args:
            model_path (str): The full path to the saved Keras model file.

        Returns:
            Model | None: The loaded Keras model, or None if loading failed.
        """
        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:
            # Ensure gaussian_nll is registered if it wasn't globally
            custom_objects = {}
            if 'gaussian_nll' in globals() or 'gaussian_nll' in locals():
                 custom_objects['gaussian_nll'] = gaussian_nll
            else:
                 # This case should ideally not happen if imports/fallbacks work
                 logger.warning("gaussian_nll custom object not found during load. Model might fail if it uses it.")

            # Load using custom_objects dictionary
            self.model = tf.keras.models.load_model(model_path, custom_objects=custom_objects)
            logger.info("GRU model loaded successfully.")
            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

    def predict(self, X_data: np.ndarray, batch_size: int = 1024) -> Any:
        """
        Generates predictions using the loaded/trained model.

        Args:
            X_data (np.ndarray): Input data sequences (shape: [n_samples, lookback, n_features]).
            batch_size (int): Batch size for prediction.

        Returns:
            Any: The model's predictions (typically a list of numpy arrays for multi-output models).
                 Returns None if no model is available or prediction fails.
        """
        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.")
            # Log shapes of predictions if it's a list (multi-output)
            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 