""" Keras-specific model implementation with TensorFlow integration.
Provides concrete implementations for Keras model operations.

Features:

- Transfer learning layer freezing/unfreezing
- Keras-specific callbacks (early stopping, LR reduction)
- Model checkpointing/weight management
- GPU-optimized prediction pipelines
- Keras metric/loss configuration
- Model serialization/deserialization

Implements ModelInterface for Keras-based models.
"""
# pyright: reportUnknownMemberType=false
# pyright: reportUnknownVariableType=false
# pyright: reportUnknownArgumentType=false
# pyright: reportArgumentType=false
# pyright: reportCallIssue=false
# pyright: reportMissingTypeStubs=false
# pyright: reportOptionalMemberAccess=false
# pyright: reportOptionalCall=false

# Imports
import gc
import multiprocessing
import multiprocessing.queues
import os
from collections.abc import Iterable
from tempfile import TemporaryDirectory
from typing import Any

import mlflow
import mlflow.keras
import numpy as np
import tensorflow as tf
from keras.backend import clear_session
from keras.callbacks import Callback, CallbackList, EarlyStopping, History, ReduceLROnPlateau, TensorBoard
from keras.layers import Dense, GlobalAveragePooling2D
from keras.losses import CategoricalCrossentropy, CategoricalFocalCrossentropy, Loss
from keras.metrics import AUC, CategoricalAccuracy, F1Score, Metric
from keras.models import Model, Sequential
from keras.optimizers import Adam, AdamW, Lion, Optimizer
from keras.utils import set_random_seed
from numpy.typing import NDArray

from ...ctx import Muffle
from ...decorators import measure_time
from ...print import colored_for_loop, debug, info, progress, warning
from .. import mlflow_utils
from ..config.get import DataScienceConfig
from ..dataset import Dataset, GroupingStrategy
from ..utils import Utils
from .keras_utils.callbacks import ColoredProgressBar, LearningRateFinder, ModelCheckpointV2, ProgressiveUnfreezing, WarmupScheduler
from .keras_utils.losses import NextGenerationLoss
from .keras_utils.visualizations import all_visualizations_for_image
from .model_interface import ModelInterface


class BaseKeras(ModelInterface):
	""" Base class for Keras models with common functionality. """

	def class_load(self) -> None:
		""" Clear the session and collect garbage, reset random seeds and call the parent class method. """
		super().class_load()
		clear_session()
		gc.collect()
		set_random_seed(DataScienceConfig.SEED)
		self.final_model: Model

	def _fit(
		self,
		model: Model,
		x: Any,
		y: Any | None = None,
		validation_data: tuple[Any, Any] | None = None,
		shuffle: bool = True,
		batch_size: int | None = None,
		epochs: int = 1,
		callbacks: list[Callback] | None = None,
		class_weight: dict[int, float] | None = None,
		verbose: int = 0,
		*args: Any,
		**kwargs: Any
	) -> History:
		""" Manually fit the model with a custom training loop instead of using model.fit().

		This method implements a custom training loop for more control over the training process.
		It's useful for implementing custom training behaviors that aren't easily done with model.fit()
		such as unfreezing layers during training, resetting the optimizer, etc.

		Args:
			model             (Model):                   The model to train
			x                 (Any):                     Training data inputs
			y                 (Any | None):              Training data targets
			validation_data   (tuple[Any, Any] | None):  Validation data as a tuple of (inputs, targets)
			shuffle           (bool):                    Whether to shuffle the training data every epoch
			batch_size        (int | None):              Number of samples per gradient update.
			epochs            (int):                     Number of epochs to train the model.
			callbacks         (list[Callback] | None):   List of callbacks to apply during training.
			class_weight      (dict[int, float] | None): Optional dictionary mapping class indices to weights.
			verbose           (int):                     Verbosity mode.

		Returns:
			History: Training history
		"""
		# Set TensorFlow to use the XLA compiler
		tf.config.optimizer.set_jit(True)

		# Build training dataset
		if y is None and isinstance(x, tf.data.Dataset):
			train_dataset: tf.data.Dataset = x
		else:
			train_dataset: tf.data.Dataset = tf.data.Dataset.from_tensor_slices((x, y) if y is not None else x)

		# Optimize dataset pipeline
		if shuffle:
			buffer_size: int = len(x) if hasattr(x, '__len__') else 10000
			buffer_size = min(buffer_size, 50000)
			train_dataset = train_dataset.shuffle(buffer_size=buffer_size, reshuffle_each_iteration=True)
		if batch_size is not None:
			train_dataset = train_dataset.batch(batch_size)
		train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE)

		# Handle validation data
		val_dataset: tf.data.Dataset | None = None
		if validation_data is not None:
			x_val, y_val = validation_data
			val_dataset = tf.data.Dataset.from_tensor_slices((x_val, y_val))
			if batch_size is not None:
				val_dataset = val_dataset.batch(batch_size)
			val_dataset = val_dataset.cache().prefetch(tf.data.AUTOTUNE)

		# Handle callbacks
		callback_list: CallbackList = CallbackList(
			callbacks,
			add_history=True,
			add_progbar=verbose != 0,
			model=model,
			verbose=verbose,
			epochs=epochs,
			steps=tf.data.experimental.cardinality(train_dataset).numpy(),
		)

		# Precompute class weights tensor outside the training loop
		class_weight_tensor: NDArray[Any] | None = None
		if class_weight:
			class_weight_values: list[float] = [float(class_weight.get(i, 1.0)) for i in range(self.num_classes)]
			class_weight_tensor = tf.constant(class_weight_values, dtype=tf.float32)

		# Precompute the gather weights function outside the training loop
		@tf.function(jit_compile=True, experimental_relax_shapes=True)
		def gather_weights(label_indices: tf.Tensor) -> tf.Tensor | None:
			if class_weight_tensor is not None:
				return tf.gather(class_weight_tensor, label_indices)
			return None

		# Get optimizer (will use loss scaling automatically under mixed-precision)
		is_ls: bool = isinstance(model.optimizer, tf.keras.mixed_precision.LossScaleOptimizer)

		# Training step with proper loss scaling
		@tf.function(jit_compile=True, experimental_relax_shapes=True)
		def train_step(xb: tf.Tensor, yb: tf.Tensor, training: bool = True) -> dict[str, Any]:
			""" Execute a single training step with gradient calculation and optimization.

			Args:
				xb (tf.Tensor): Input batch data
				yb (tf.Tensor): Target batch data

			Returns:
				dict[str, Any]: The metrics for the training step
			"""
			labels = tf.cast(tf.argmax(yb, axis=1), tf.int32)
			sw = gather_weights(labels)
			with tf.GradientTape(watch_accessed_variables=training) as tape:
				preds = model(xb, training=training)
				loss = model.compiled_loss(yb, preds, sample_weight=sw)
				loss = tf.reduce_mean(loss)

				# Scale loss if using LossScaleOptimizer
				if is_ls:
					loss = model.optimizer.get_scaled_loss(loss)

			# Backpropagate the loss
			if training:
				model.optimizer.minimize(loss, model.trainable_weights, tape=tape)

			# Update the metrics
			model.compiled_metrics.update_state(yb, preds, sample_weight=sw)
			return model.get_metrics_result()

		# Start callbacks
		logs: dict[str, Any] = {"loss": 0.0}
		callback_list.on_train_begin()

		# Custom training loop
		for epoch in range(epochs):

			# Callbacks and reset metrics
			callback_list.on_epoch_begin(epoch)
			model.compiled_metrics.reset_state()
			model.compiled_loss.reset_state()

			# Train on all batches
			for step, (x_batch, y_batch) in enumerate(train_dataset):
				callback_list.on_batch_begin(step)
				logs.update(train_step(x_batch, y_batch, training=True))
				callback_list.on_batch_end(step, logs)

			# Compute metrics for validation
			if val_dataset is not None:
				model.compiled_metrics.reset_state()
				model.compiled_loss.reset_state()

				# Run through all validation data
				for x_val, y_val in val_dataset:
					train_step(x_val, y_val, training=False)

				# Prefix "val_" to the metrics
				for key, value in model.get_metrics_result().items():
					logs[f"val_{key}"] = value

			callback_list.on_epoch_end(epoch, logs)
		callback_list.on_train_end(logs)

		# Return history
		return model.history # pyright: ignore [reportReturnType]


	def _get_architectures(
		self, optimizer: Any = None, loss: Any = None, metrics: list[Any] | None = None
	) -> tuple[Model, Model]:
		""" Get the model architecture and compile it if enough information is provided.

		This method builds and returns the model architecture.
		If optimizer, loss, and (optionally) metrics are provided, the model will be compiled.

		Args:
			optimizer (Any): The optimizer to use for training
			loss (Any): The loss function to use for training
			metrics (list[Any] | None): The metrics to use for evaluation
		Returns:
			tuple[Model, Model]: The final model and the base model
		"""

		# Get the base model (use imagenet anyway)
		base_model: Model = self._get_base_model()

		# Add a top layer since the base model doesn't have one
		output_layer: Model = Sequential([
			GlobalAveragePooling2D(),
			Dense(self.num_classes, activation="softmax")
		])(base_model.output)
		final_model: Model = Model(inputs=base_model.input, outputs=output_layer)

		# If no optimizer is provided, return the uncompiled models
		if optimizer is None:
			return final_model, base_model

		# Load transfer learning weights if provided
		if os.path.exists(self.transfer_learning):
			try:
				final_model.load_weights(self.transfer_learning)
				info(f"Transfer learning weights loaded from '{self.transfer_learning}'")
			except Exception as e:
				warning(f"Error loading transfer learning weights from '{self.transfer_learning}': {e}")

		# Freeze the base model except for the last layers (if unfreeze percentage is less than 100%)
		if self.unfreeze_percentage < 100.0:
			base_model.trainable = False
			last_layers: list[Model] = base_model.layers[-self.fine_tune_last_layers:]
			for layer in last_layers:
				layer.trainable = True
			info(
				f"Fine-tune from layer {max(0, len(base_model.layers) - self.fine_tune_last_layers)} "
				f"to {len(base_model.layers)} ({self.fine_tune_last_layers} layers)"
			)

		# Add XLA specific optimizations for compilation
		compile_options = {}
		if hasattr(tf.config.optimizer, "get_jit") and tf.config.optimizer.get_jit():
			compile_options["steps_per_execution"] = 10  # Batch multiple steps for XLA

		# Compile the model and return it
		final_model.compile(
			optimizer=optimizer,
			loss=loss,
			metrics=metrics if metrics is not None else [],
			jit_compile=True,
			**compile_options
		)
		return final_model, base_model


	# Protected methods for training
	def _get_callbacks(self) -> list[Callback]:
		""" Get the callbacks for training. """
		callbacks: list[Callback] = []

		# Add warmup scheduler if enabled
		if self.warmup_epochs > 0:
			warmup_scheduler: WarmupScheduler = WarmupScheduler(
				warmup_epochs=self.warmup_epochs,
				initial_lr=self.initial_warmup_lr,
				target_lr=self.learning_rate
			)
			callbacks.append(warmup_scheduler)

		# Add ReduceLROnPlateau
		callbacks.append(ReduceLROnPlateau(
			monitor="val_loss",
			mode="min",
			factor=self.factor,
			patience=self.reduce_lr_patience,
			min_delta=self.min_delta,
			min_lr=self.min_lr
		))

		# Add TensorBoard for profiling
		log_dir: str = f"{DataScienceConfig.TENSORBOARD_FOLDER}/{self.run_name}"
		os.makedirs(log_dir, exist_ok=True)
		callbacks.append(TensorBoard(
			log_dir=log_dir,
			histogram_freq=1,  # Log histogram visualizations every epoch
			profile_batch=(10, 20)  # Profile batches 10-20
		))

		# Add EarlyStopping to prevent overfitting
		callbacks.append(EarlyStopping(
			monitor="val_loss",
			mode="min",
			patience=self.early_stop_patience,
			verbose=0
		))
		return callbacks

	def _get_metrics(self) -> list[Metric]:
		""" Get the metrics for training.

		Returns:
			list: List of metrics to track during training including accuracy, AUC, etc.
		"""
		# Fix the F1Score dtype if mixed precision is enabled
		f1score_dtype: tf.DType = tf.float16 if DataScienceConfig.MIXED_PRECISION_POLICY == "mixed_float16" else tf.float32
		f1score: F1Score = F1Score(name="f1_score", average="macro", dtype=f1score_dtype)
		f1score.beta = tf.constant(1.0, dtype=f1score_dtype) # pyright: ignore [reportAttributeAccessIssue]

		return [
			CategoricalAccuracy(name="categorical_accuracy"),
			AUC(name="auc"),
			f1score,
		]

	def _get_optimizer(self, learning_rate: float = 0.0, mode: int = 1) -> Optimizer:
		""" Get the optimizer for training.

		Args:
			learning_rate  (float):  Learning rate
			mode           (int):    Mode to use
		Returns:
			Optimizer: Optimizer
		"""
		lr: float = self.learning_rate if learning_rate == 0.0 else learning_rate
		if mode == 0:
			return Adam(lr, self.beta_1, self.beta_2)
		elif mode == 1:
			return AdamW(lr, self.beta_1, self.beta_2)
		else:
			return Lion(lr)

	def _get_loss(self, mode: int = 0) -> Loss:
		""" Get the loss function for training depending on the mode.

		- 0: CategoricalCrossentropy (default)
		- 1: CategoricalFocalCrossentropy
		- 2: Next Generation Loss (with alpha = 2.4092)

		Args:
			mode (int): Mode to use
		Returns:
			Loss: Loss function
		"""
		if mode == 0:
			return CategoricalCrossentropy(name="categorical_crossentropy")
		elif mode == 1:
			return CategoricalFocalCrossentropy(name="categorical_focal_crossentropy")
		elif mode == 2:
			return NextGenerationLoss(name="ngl_loss")
		else:
			raise ValueError(f"Invalid mode: {mode}")

	def _find_best_learning_rate_subprocess(
		self, dataset: Dataset, queue: multiprocessing.queues.Queue | None = None, verbose: int = 0 # type: ignore
	) -> dict[str, Any] | None:
		""" Helper to run learning rate finder, potentially in a subprocess.

		Args:
			dataset         (Dataset):                         Dataset to use for training.
			queue           (multiprocessing.Queue | None):    Queue to put results in (if running in subprocess).
			verbose         (int):                             Verbosity level.

		Returns:
			dict[str, Any] | None: Return values
		"""
		X_train, y_train, _ = (dataset.training_data + self.additional_training_data).ungrouped_array()

		# Set random seeds for reproducibility within the process/subprocess
		set_random_seed(DataScienceConfig.SEED)

		# Create LR finder callback
		lr_finder: LearningRateFinder = LearningRateFinder(
			min_lr=self.lr_finder_min_lr,
			max_lr=self.lr_finder_max_lr,
			steps_per_epoch=np.ceil(len(X_train) / self.batch_size),
			epochs=self.lr_finder_epochs,
			update_per_epoch=self.lr_finder_update_per_epoch,
			update_interval=self.lr_finder_update_interval
		)

		# Get compiled model with the optimizer and loss
		final_model, _ = self._get_architectures(self._get_optimizer(), self._get_loss())

		# Create callbacks
		callbacks: list[Callback] = [lr_finder]
		if verbose > 0:
			callbacks.append(ColoredProgressBar("LR Finder", show_lr=True))

		# Run a mini training to find the best learning rate
		self._fit(
			final_model,
			X_train, y_train,
			batch_size=self.batch_size,
			epochs=self.lr_finder_epochs,
			callbacks=callbacks,
			class_weight=self.class_weight,
			verbose=0
		)

		# Prepare results
		results: dict[str, Any] = {
			"learning_rates": lr_finder.learning_rates,
			"losses": lr_finder.losses
		}

		# Return values if no queue, otherwise put them in the queue
		if queue is None:
			return results
		else:
			return queue.put(results)

	def _find_best_unfreeze_percentage_subprocess(
		self, dataset: Dataset, queue: multiprocessing.queues.Queue | None = None, verbose: int = 0 # type: ignore
	) -> dict[str, Any] | None:
		""" Helper to run unfreeze percentage finder, potentially in a subprocess.

		Args:
			dataset         (Dataset):                         Dataset to use for training.
			queue           (multiprocessing.Queue | None):    Queue to put results in (if running in subprocess).
			verbose         (int):                             Verbosity level.

		Returns:
			dict[str, Any] | None: Return values
		"""
		X_train, y_train, _ = (dataset.training_data + self.additional_training_data).ungrouped_array()

		# Set random seeds for reproducibility within the process/subprocess
		set_random_seed(DataScienceConfig.SEED)

		# Get compiled model with the optimizer and loss
		lr: float = self.learning_rate
		optimizer = self._get_optimizer(lr)
		loss_fn = self._get_loss()
		final_model, base_model = self._get_architectures(optimizer, loss_fn)

		# Function to get compiled optimizer
		def get_compiled_optimizer() -> Optimizer:
			optimizer: Optimizer = self._get_optimizer(lr)
			return final_model._get_optimizer(optimizer) # pyright: ignore [reportPrivateUsage]

		# Create unfreeze finder callback
		unfreeze_finder: ProgressiveUnfreezing = ProgressiveUnfreezing(
			base_model=base_model,
			steps_per_epoch=np.ceil(len(X_train) / self.batch_size),
			epochs=self.unfreeze_finder_epochs,
			reset_weights=True,
			reset_optimizer_function=get_compiled_optimizer,
			update_per_epoch=self.unfreeze_finder_update_per_epoch,
			update_interval=self.unfreeze_finder_update_interval,
			progressive_freeze=True		# Start from 100% unfrozen to 0% unfrozen to prevent biases
		)

		# Create callbacks
		callbacks: list[Callback] = [unfreeze_finder]
		if verbose > 0:
			callbacks.append(ColoredProgressBar("Unfreeze Finder"))

		self._fit(
			final_model,
			X_train, y_train,
			batch_size=self.batch_size,
			epochs=self.unfreeze_finder_epochs,
			callbacks=callbacks,
			class_weight=self.class_weight,
			verbose=0
		)

		# Prepare results
		unfreeze_percentages, losses = unfreeze_finder.get_results()
		results: dict[str, Any] = {
			"unfreeze_percentages": unfreeze_percentages,
			"losses": losses
		}

		# Return values if no queue, otherwise put them in the queue
		if queue is None:
			return results
		else:
			return queue.put(results)

	def _train_subprocess(
		self,
		dataset: Dataset,
		checkpoint_path: str,
		temp_dir: TemporaryDirectory[str] | None = None,
		queue: multiprocessing.queues.Queue | None = None, # type: ignore
		verbose: int = 0
	) -> dict[str, Any] | None:
		""" Train the model in a subprocess.

		The reason for this is that when training too much models on the same process,
		your process may be killed by the OS since it used too much resources over time.
		So we train each model in a separate process to avoid this issue.

		Args:
			model            (Model):                           Model to train
			dataset          (Dataset):                         Dataset to train on
			checkpoint_path  (str):                             Path to save the best model checkpoint
			temp_dir         (TemporaryDirectory[str] | None):  Temporary directory to save the visualizations
			queue            (multiprocessing.Queue | None):    Queue to put the history in
			verbose          (int):                             Verbosity level
		Returns:
			dict[str, Any]: Return values
		"""
		to_return: dict[str, Any] = {}
		set_random_seed(DataScienceConfig.SEED)

		# Extract the training and validation data
		X_train, y_train, _ = (dataset.training_data + self.additional_training_data).ungrouped_array()
		X_val, y_val, _ = dataset.val_data.ungrouped_array()
		X_test, y_test, test_filepaths = dataset.test_data.ungrouped_array()
		true_classes: NDArray[Any] = Utils.convert_to_class_indices(y_val)

		# Create the checkpoint callback
		os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
		model_checkpoint: ModelCheckpointV2 = ModelCheckpointV2(
			epochs_before_start=self.model_checkpoint_delay,
			filepath=checkpoint_path,
			monitor="val_loss",
			mode="min",
			save_best_only=True,
			save_weights_only=True,
			verbose=0
		)

		# Get the compiled model
		model, _ = self._get_architectures(self._get_optimizer(), self._get_loss(), self._get_metrics())

		# Create the callbacks, add the progress bar if verbose is 1
		callbacks = [model_checkpoint, *self._get_callbacks()]
		if verbose > 0:
			callbacks.append(ColoredProgressBar("Training", show_lr=True))

		# Train the model
		history: History = self._fit(
			model,
			X_train, y_train,
			validation_data=(X_val, y_val),
			batch_size=self.batch_size,
			epochs=self.epochs,
			callbacks=callbacks,
			class_weight=self.class_weight,
			verbose=0
		)

		# Load the best model from the checkpoint file and remove it
		debug(f"Loading best model from '{checkpoint_path}'")
		model.load_weights(checkpoint_path)
		os.remove(checkpoint_path)
		debug(f"Best model loaded from '{checkpoint_path}', deleting it...")

		# Evaluate the model
		to_return["history"] = history.history
		to_return["eval_results"] = model.evaluate(X_test, y_test, return_dict=True, verbose=0)
		to_return["predictions"] = model.predict(X_test, verbose=0)
		to_return["true_classes"] = true_classes
		to_return["training_predictions"] = model.predict(X_train, verbose=0)
		to_return["training_true_classes"] = Utils.convert_to_class_indices(y_train)

		# --- Visualization Generation (Using viz_kwargs) ---
		if temp_dir is not None:

			# Ensure fold_number > 0 for LOO visualization
			test_images: list[NDArray[Any]] = list(X_test)

			# Prepare the arguments for the visualizations
			viz_args_list: list[tuple[NDArray[Any], Any, tuple[str, ...], str]] = [
				(test_images[i], true_classes[i], test_filepaths[i], "test_folds")
				for i in range(len(test_images))
			]

			# Generate visualizations in the provided temporary directory
			for img_viz, label_idx, files, data_type in viz_args_list:
				# Extract the base name of the file/group
				if dataset.grouping_strategy == GroupingStrategy.NONE:
					base_name: str = os.path.splitext(os.path.basename(files[0]))[0]
				else:
					base_name: str = os.path.basename(os.path.dirname(files[0]))

				# Generate all visualizations for the image
				all_visualizations_for_image(
					model=model, # Use the trained model from this subprocess
					folder_path=temp_dir.name,
					img=img_viz,
					base_name=base_name,
					class_idx=label_idx,
					class_name=dataset.labels[label_idx],
					files=files,
					data_type=data_type,
				)

		# Return values if no queue, otherwise put them in the queue
		if queue is None:
			to_return["model"] = model	# Add the trained model to the return values if not in a subprocess
			return to_return
		else:
			return queue.put(to_return)


	# Predict method
	def class_predict(self, X_test: Iterable[NDArray[Any]] | tf.data.Dataset) -> Iterable[NDArray[Any]]:
		""" Predict the class for the given input data.

		Args:
			X_test  (Iterable[NDArray[Any]]):  List of inputs to predict (e.g. a batch of images)
		Returns:
			Iterable[NDArray[Any]]: A batch of predictions (model.predict())
		"""
		# Create a tf.data.Dataset to avoid retracing
		if isinstance(X_test, tf.data.Dataset):
			dataset: tf.data.Dataset = X_test
			was_dataset: bool = True
		else:
			dataset: tf.data.Dataset = tf.data.Dataset.from_tensor_slices(X_test).batch(32).prefetch(tf.data.AUTOTUNE)
			was_dataset: bool = False

		# Create an optimized prediction function
		@tf.function(jit_compile=True)
		def optimized_predict(x_batch: tf.Tensor) -> tf.Tensor:
			return self.final_model(x_batch, training=False)

		# For each model, predict the class
		model_preds: list[NDArray[Any]] = []
		for batch in dataset:
			pred: tf.Tensor = optimized_predict(batch)
			model_preds.append(pred.numpy())

		# Clear RAM
		if not was_dataset:
			del dataset
		gc.collect()

		# Return the predictions
		return np.concatenate(model_preds) if model_preds else np.array([])


	# Protected methods for evaluation
	@measure_time(progress)
	def _log_final_model(self) -> None:
		""" Log the best model (and its weights). """
		with Muffle(mute_stderr=True):
			mlflow.keras.log_model(self.final_model, "best_model") # pyright: ignore [reportPrivateImportUsage]
		mlflow.set_tag(key="has_saved_model", value="True")

		# Get the weights path and create the directory if it doesn't exist
		weights_path: str = mlflow_utils.get_weights_path()
		os.makedirs(os.path.dirname(weights_path), exist_ok=True)

		# Save the best model's weights without the last layer
		self.final_model.save_weights(weights_path)


	def class_evaluate(
		self, dataset: Dataset, metrics_names: tuple[str, ...] = (), save_model: bool = False, verbose: int = 0
	) -> bool:
		""" Evaluate the model using the given predictions and labels.

		Args:
			dataset         (Dataset):   Dataset containing the training and testing data
			metrics_names   (list[str]): List of metrics to plot (default to all metrics)
			save_model      (bool):      Whether to save the best model
			verbose         (int):       Level of verbosity
		Returns:
			bool: True if evaluation was successful
		"""
		# First perform standard evaluation from parent class
		result: bool = super().class_evaluate(dataset, metrics_names, save_model, verbose)
		if not DataScienceConfig.DO_SALIENCY_AND_GRADCAM:
			return result

		# Get test and train data
		X_test, y_test, test_filepaths = dataset.test_data.ungrouped_array()
		test_images: list[NDArray[Any]] = list(X_test)
		test_labels: list[int] = Utils.convert_to_class_indices(y_test).tolist()

		X_train, y_train, train_filepaths = dataset.training_data.remove_augmented_files().ungrouped_array()
		train_images: list[NDArray[Any]] = list(X_train)
		train_labels: list[int] = Utils.convert_to_class_indices(y_train).tolist()

		# Process test images
		test_args_list: list[tuple[NDArray[Any], int, tuple[str, ...], str]] = [
			(test_images[i], test_labels[i], test_filepaths[i], "test")
			for i in range(min(100, len(test_images)))
		]

		# Process train images
		train_args_list: list[tuple[NDArray[Any], int, tuple[str, ...], str]] = [
			(train_images[i], train_labels[i], train_filepaths[i], "train")
			for i in range(min(10, len(train_images)))
		]

		# Combine both lists
		all_args_list = test_args_list + train_args_list

		# Create the description
		desc: str = ""
		if verbose > 0:
			desc = f"Generating visualizations for {len(test_args_list)} test and {len(train_args_list)} train images"

		# For each image, generate all visualizations, then log them to MLFlow
		with TemporaryDirectory() as temp_dir:
			for img, label, files, data_type in colored_for_loop(all_args_list, desc=desc):

				# Extract the base name of the file
				if dataset.grouping_strategy == GroupingStrategy.NONE:
					base_name: str = os.path.splitext(os.path.basename(files[0]))[0]
				else:
					base_name: str = os.path.basename(os.path.dirname(files[0]))

				# Generate all visualizations for the image
				all_visualizations_for_image(
					model=self.final_model,
					folder_path=temp_dir,
					img=img,
					base_name=base_name,
					class_idx=label,
					class_name=dataset.labels[label],
					files=files,
					data_type=data_type,
				)

			# Log the visualizations
			mlflow.log_artifacts(temp_dir)

		return result

