"""
GRU Hyperparameter Tuner.

Implements hyperparameter optimization for GRU models using Optuna.
"""

import os
import logging
import json
from typing import Dict, Any, List, Tuple, Optional
import numpy as np
import pandas as pd
import optuna
from optuna.pruners import MedianPruner
from optuna.samplers import TPESampler
import tensorflow as tf
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
from tensorflow.keras.callbacks import Callback

from gru_sac_predictor.src.gru_model_handler import GRUModelHandler
from gru_sac_predictor.src.metrics import edge_filtered_accuracy, calculate_brier_score

logger = logging.getLogger(__name__)

# --- Revision 3: Optuna Pruning Callback --- #
class OptunaPruningCallback(Callback):
    """Keras Callback for Optuna Pruning."""
    def __init__(self, trial: optuna.Trial, monitor: str = 'val_loss'):
        super().__init__()
        self._trial = trial
        self._monitor = monitor

    def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, float]] = None):
        logs = logs or {}
        current_value = logs.get(self._monitor)
        if current_value is None:
            # Metric not found, cannot prune based on it
            # logger.warning(f"Optuna Callback: Metric '{self._monitor}' not found in logs. Skipping pruning check.")
            return

        # Report current value to Optuna trial
        self._trial.report(current_value, step=epoch)

        # Check if the trial should be pruned
        if self._trial.should_prune():
            message = f"Trial pruned at epoch {epoch} based on {self._monitor}={current_value}."
            logger.info(message)
            # Stop training by raising TrialPruned exception
            raise optuna.TrialPruned(message)
# --- End Revision 3 --- #

class GRUHyperTuner:
    """
    Optimizes GRU model hyperparameters using Optuna.
    
    This class runs an Optuna hyperparameter sweep to find the best GRU 
    hyperparameters for a given dataset using edge-filtered accuracy as
    the objective function.
    """
    
    def __init__(self, config: Dict[str, Any], fold_dir: str):
        """
        Initialize the GRU Hyperparameter Tuner.
        
        Args:
            config (Dict[str, Any]): Configuration dictionary.
            fold_dir (str): Directory to store fold-specific results.
        """
        self.config = config
        self.fold_dir = fold_dir
        self.is_ternary = config.get('gru', {}).get('use_ternary', False)
        
        # Get hyperparameter sweep config
        sweep_config = config.get('hyperparameter_tuning', {}).get('gru', {})
        self.n_trials = sweep_config.get('sweep_n_trials', 20)
        self.timeout = sweep_config.get('sweep_timeout', 7200)  # 2 hours default
        self.edge_threshold = config.get('calibration', {}).get('edge_threshold', 0.1)
        
        # Pruning settings
        self.pruning_enabled = sweep_config.get('enable_pruning', True)
        self.pruning_warmup = sweep_config.get('pruning_warmup_trials', 5)
        
        # Create fold's best params directory if it doesn't exist
        os.makedirs(fold_dir, exist_ok=True)
        
        # Set random seeds for reproducibility
        self.seed = config.get('random_seed', 42)
        np.random.seed(self.seed)
        tf.random.set_seed(self.seed)
        
        # Store best params and score
        self.best_params = None
        self.best_score = 0.0
        self.best_trial = None
        
    def _create_model(self, trial: optuna.Trial, X_train_shape: Tuple) -> Tuple[GRUModelHandler, int, int]:
        """
        Create a GRU model with hyperparameters suggested by Optuna trial.
        
        Args:
            trial (optuna.Trial): Optuna trial object.
            X_train_shape (Tuple): Shape of X_train to determine input dimensions.
            
        Returns:
            Tuple[GRUModelHandler, int, int]: Configured GRU model handler, lookback, and n_features.
        """
        # Create a unique run ID for this trial
        trial_run_id = f"tune_trial_{trial.number}"
        
        # Create a copy of the config to modify
        trial_config = dict(self.config)
        
        # Sample hyperparameters
        gru_config = trial_config.get('gru', {})
        
        # Units in GRU layers
        gru_config['units1'] = trial.suggest_categorical('units1', [48, 64, 96])
        gru_config['units2'] = trial.suggest_categorical('units2', [0, 48, 64])  # 0 means no second layer
        
        # Dropout rates
        gru_config['dropout1'] = trial.suggest_float('dropout1', 0.0, 0.5)
        gru_config['dropout2'] = trial.suggest_float('dropout2', 0.0, 0.5)
        
        # Learning rate and batch size
        gru_config['learning_rate'] = trial.suggest_float('learning_rate', 1e-4, 3e-3, log=True)
        gru_config['batch_size'] = trial.suggest_categorical('batch_size', [32, 64, 128, 256])
        
        # L2 regularization
        gru_config['l2_reg'] = trial.suggest_float('l2_reg', 1e-6, 1e-3, log=True)
        
        # Add attention heads if using v3 model
        if trial_config.get('control', {}).get('use_v3', False):
            gru_config['attn_heads'] = trial.suggest_categorical('attn_heads', [0, 4]) # 0 implies None
        
        # Update config
        trial_config['gru'] = gru_config
        
        # Create model handler with trial config
        model_handler = GRUModelHandler(
            run_id=trial_run_id,
            models_dir=self.fold_dir,
            config=trial_config
        )
        
        # Set input shape and configuration
        lookback = gru_config.get('lookback', 60)
        n_features = X_train_shape[2]
        
        return model_handler, lookback, n_features
    
    def objective(self, trial: optuna.Trial, X_train: np.ndarray, y_train_dict: Dict[str, np.ndarray],
                 X_val: np.ndarray, y_val_dict: Dict[str, np.ndarray]) -> float:
        """
        Objective function for Optuna optimization.
        
        Args:
            trial (optuna.Trial): Optuna trial object.
            X_train (np.ndarray): Training features.
            y_train_dict (Dict[str, np.ndarray]): Training targets.
            X_val (np.ndarray): Validation features.
            y_val_dict (Dict[str, np.ndarray]): Validation targets.
            
        Returns:
            float: Edge-filtered accuracy score to maximize.
        """
        # Create model with trial hyperparameters
        model_handler, lookback, n_features = self._create_model(trial, X_train.shape)
        
        # Early stopping config
        patience = self.config.get('gru', {}).get('patience', 5)
        max_epochs = self.config.get('gru', {}).get('epochs', 25)
        batch_size = model_handler.config.get('gru', {}).get('batch_size', 128)
        
        # --- Revision 3: Add Optuna Pruning Callback --- #
        callbacks = []
        if self.pruning_enabled:
            # Use val_loss as the monitor for pruning
            pruning_callback = OptunaPruningCallback(trial, monitor='val_loss')
            callbacks.append(pruning_callback)
        # --- End Revision 3 --- #
        
        # Determine which target key to use based on ternary flag
        dir_key = "dir3" if self.is_ternary else "dir"
        
        # Train the model
        try:
            gru_model, history = model_handler.train(
                X_train=X_train,
                y_train_dict=y_train_dict,
                X_val=X_val,
                y_val_dict=y_val_dict,
                lookback=lookback,
                n_features=n_features,
                max_epochs=max_epochs,
                batch_size=batch_size,
                patience=patience,
                callbacks=callbacks # Pass the Optuna callback
            )
            
            if gru_model is None:
                logger.warning(f"Trial {trial.number}: Model training failed.")
                return -np.inf # Return worse possible score for maximization
                
            # Calculate edge-filtered accuracy and Brier score
            edge_accuracy = np.nan
            brier_score = np.nan
            n_filtered = 0
            objective_value = -np.inf # Default for maximization objective
            
            if self.is_ternary:
                # Get logits and calibrate if needed
                logits_val = model_handler.predict_logits(X_val)
                if logits_val is None:
                    logger.warning(f"Trial {trial.number}: Failed to get logits for ternary evaluation.")
                    return -np.inf
                    
                # Use P(up) equivalent for binary check compatibility
                p_up_equiv = logits_val[:, 2]
                y_true_binary_equiv = (np.argmax(y_val_dict[dir_key], axis=1) == 2).astype(int)
                
                edge_accuracy, n_filtered = edge_filtered_accuracy(
                    y_true=y_true_binary_equiv,
                    p_cal=p_up_equiv,
                    thr=self.edge_threshold
                )
                
                # Brier score not calculated for ternary in this setup
                logger.warning("Brier score calculation skipped for ternary objective.")
                brier_score = 0.0 # Assign neutral value for objective calculation
                objective_value = edge_accuracy # Objective for ternary is just edge accuracy for now
                
            else:
                # Get binary predictions
                predictions_val = model_handler.predict(X_val)
                if predictions_val is None or len(predictions_val) < 3:
                    logger.warning(f"Trial {trial.number}: Failed to get predictions for binary evaluation.")
                    return -np.inf
                    
                p_cal = predictions_val[2].flatten()
                y_true = y_val_dict[dir_key]
                
                edge_accuracy, n_filtered = edge_filtered_accuracy(
                    y_true=y_true,
                    p_cal=p_cal,
                    thr=self.edge_threshold
                )
                
                # Calculate Brier score
                try:
                    brier_score = calculate_brier_score(y_true, p_cal)
                except Exception as e:
                    logger.warning(f"Trial {trial.number}: Failed to calculate Brier score: {e}. Setting to NaN.")
                    brier_score = np.nan

            # --- Calculate Combined Objective --- #
            objective_metric = self.config.get('hyperparameter_tuning', {}).get('gru', {}).get('objective_metric', 'edge_accuracy')
            if pd.isna(edge_accuracy) or pd.isna(brier_score):
                objective_value = -np.inf # Invalid if metrics are NaN
            elif objective_metric == 'edge_accuracy':
                objective_value = edge_accuracy
            elif objective_metric == 'brier_score':
                objective_value = -brier_score # Minimize Brier, so maximize negative Brier
            elif objective_metric == 'edge_acc_minus_brier':
                w_edge = self.config.get('hyperparameter_tuning', {}).get('gru', {}).get('objective_edge_acc_weight', 0.7)
                w_brier = self.config.get('hyperparameter_tuning', {}).get('gru', {}).get('objective_brier_weight', 0.3)
                objective_value = w_edge * edge_accuracy - w_brier * brier_score
            else:
                logger.warning(f"Unknown objective metric '{objective_metric}'. Defaulting to edge accuracy.")
                objective_value = edge_accuracy
            # --- End Combined Objective --- #

            # --- Log components to Optuna Trial --- #
            trial.set_user_attr("edge_accuracy", float(edge_accuracy) if pd.notna(edge_accuracy) else None)
            trial.set_user_attr("brier_score", float(brier_score) if pd.notna(brier_score) else None)
            trial.set_user_attr("n_filtered", int(n_filtered))
            # --- End Log --- #

            logger.info(f"Trial {trial.number} finished. Objective ({objective_metric}): {objective_value:.4f} (EdgeAcc: {edge_accuracy:.4f}, Brier: {brier_score:.4f}, N_filt: {n_filtered})")
            return objective_value
            
        except optuna.TrialPruned as e:
            # Handle pruning gracefully
            logger.info(f"Trial {trial.number} pruned: {e}")
            raise e # Re-raise to let Optuna handle it
        except Exception as e:
            logger.error(f"Error during Optuna objective evaluation (Trial {trial.number}): {e}", exc_info=True)
            # Return a very bad value for maximization problems
            return -np.inf
            
    def optimize(self, X_train: np.ndarray, y_train_dict: Dict[str, np.ndarray],
                X_val: np.ndarray, y_val_dict: Dict[str, np.ndarray]) -> Dict[str, Any]:
        """
        Run Optuna optimization to find the best hyperparameters.
        
        Args:
            X_train (np.ndarray): Training features.
            y_train_dict (Dict[str, np.ndarray]): Training targets.
            X_val (np.ndarray): Validation features.
            y_val_dict (Dict[str, np.ndarray]): Validation targets.
            
        Returns:
            Dict[str, Any]: Best hyperparameters.
        """
        # Create sampler and pruner
        sampler = TPESampler(seed=self.seed)
        pruner = MedianPruner(n_startup_trials=self.pruning_warmup, n_warmup_steps=0) if self.pruning_enabled else None
        
        # Create study
        study = optuna.create_study(
            direction="maximize",
            sampler=sampler,
            pruner=pruner,
            study_name="gru_hyperparameter_tuning"
        )
        
        # Define objective function with fixed data
        objective_with_data = lambda trial: self.objective(
            trial, X_train, y_train_dict, X_val, y_val_dict
        )
        
        # Run optimization
        logger.info(f"Starting GRU hyperparameter optimization with {self.n_trials} trials")
        study.optimize(objective_with_data, n_trials=self.n_trials, timeout=self.timeout)
        
        # Get best trial and params
        self.best_trial = study.best_trial
        self.best_params = study.best_params
        self.best_score = study.best_value
        
        # Log results
        logger.info(f"Optimization finished. Best score: {self.best_score:.4f}")
        logger.info(f"Best parameters: {self.best_params}")
        
        # Save best parameters to JSON file
        best_params_path = os.path.join(self.fold_dir, 'best_gru_params.json')
        with open(best_params_path, 'w') as f:
            json.dump(self.best_params, f, indent=4)
        logger.info(f"Best parameters saved to {best_params_path}")
        
        # Plot optimization history
        self._plot_optimization_history(study)
        
        return self.best_params
        
    def _plot_optimization_history(self, study: optuna.Study) -> None:
        """
        Plot the optimization history.
        
        Args:
            study (optuna.Study): Completed Optuna study.
        """
        try:
            # Create figure
            fig = plt.figure(figsize=(10, 6))
            
            # Plot optimization history
            optuna.visualization.matplotlib.plot_optimization_history(study)
            
            # Save figure
            history_path = os.path.join(self.fold_dir, 'optuna_history.png')
            plt.savefig(history_path, dpi=100)
            plt.close(fig)
            logger.info(f"Optimization history plot saved to {history_path}")
            
            # Plot parameter importances if we have enough trials
            if len(study.trials) >= 5:
                fig = plt.figure(figsize=(10, 6))
                optuna.visualization.matplotlib.plot_param_importances(study)
                importances_path = os.path.join(self.fold_dir, 'param_importances.png')
                plt.savefig(importances_path, dpi=100)
                plt.close(fig)
                logger.info(f"Parameter importances plot saved to {importances_path}")
                
        except Exception as e:
            logger.error(f"Failed to plot optimization results: {str(e)}")
            
    def train_with_best_params(self, X_train: np.ndarray, y_train_dict: Dict[str, np.ndarray],
                              X_val: np.ndarray, y_val_dict: Dict[str, np.ndarray]) -> Tuple[GRUModelHandler, Dict]:
        """
        Train final model with the best hyperparameters.
        
        Args:
            X_train (np.ndarray): Training features.
            y_train_dict (Dict[str, np.ndarray]): Training targets.
            X_val (np.ndarray): Validation features.
            y_val_dict (Dict[str, np.ndarray]): Validation targets.
            
        Returns:
            Tuple[GRUModelHandler, Dict]: Trained model handler and history.
        """
        if self.best_params is None:
            logger.error("No best parameters found. Run optimize() first.")
            return None, None
            
        # Create a unique run ID for the final model
        final_run_id = f"tune_best_model"
        
        # Create a copy of the config to modify
        final_config = dict(self.config)
        gru_config = final_config.get('gru', {}).copy()
        
        # Apply best parameters
        for param, value in self.best_params.items():
            gru_config[param] = value
            
        # Skip second GRU layer if units2 is 0
        if 'units2' in self.best_params and self.best_params['units2'] == 0:
            gru_config['use_second_layer'] = False
        else:
            gru_config['use_second_layer'] = True
            
        # Update config
        final_config['gru'] = gru_config
        
        # Create model handler with final config
        model_handler = GRUModelHandler(
            run_id=final_run_id,
            models_dir=self.fold_dir,
            config=final_config
        )
        
        # Set training parameters
        lookback = gru_config.get('lookback', 60)
        n_features = X_train.shape[2]
        max_epochs = gru_config.get('epochs', 25)
        batch_size = gru_config.get('batch_size', 128)
        patience = gru_config.get('patience', 5)
        
        # Train final model
        logger.info("Training final model with best hyperparameters")
        model, history = model_handler.train(
            X_train=X_train,
            y_train_dict=y_train_dict,
            X_val=X_val,
            y_val_dict=y_val_dict,
            lookback=lookback,
            n_features=n_features,
            max_epochs=max_epochs,
            batch_size=batch_size,
            patience=patience
        )
        
        # Save the model
        if model is not None:
            saved_path = model_handler.save()
            logger.info(f"Final model saved to {saved_path}")
            
        return model_handler, history 