sientia-dataops-model-manager/model_manager/utils/repository/model_repository.py

"""
Model Monitoring Repository

This module contains the ModelMonitoringRepository class,
which is responsible for handling the communication with the Model Monitoring API.

It includes the methods that are used to answer ModelMonitoringService
requests using the Model Monitoring API functions.

By Monitoring we mean the evaluation of the performance of models, the generation of reports.

"""

import shutil
import traceback
from datetime import datetime
from os import makedirs, path, remove

import mlflow
import numpy as np
import pandas as pd
from sientia.ModelServing import ModelServing  # type: ignore[import-untyped]
from sientia.reports import Reports  # type: ignore[import-untyped]
from sientia_do.observability.logger import Logger
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ

from model_manager.utils.models.train_model_result import TrainModelResult


class MLFlowRepository:
    def __init__(self, host, username, password, logger: Logger):
        self.model_serving = ModelServing(
            tracking_uri=host, username=username, password=password, logger=logger
        )
        self.logger = logger

    def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
        """
        Detect and parse datetime index from data. index must be a timestamp like column.
        This function must detect the timestamp type (pandas Timestamp or datetime) and convert it to DATETIME_FORMAT_WITH_TZ.
        If the index is a string, must be in format DATETIME_FORMAT_WITH_TZ.
        If another type or format, must raise an error.
        """
        index = data.index

        # Get type of first element of index
        index_type = type(index[0])

        self.logger.custom_info(f'Index type: {index_type}', metadata)

        message = f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}'

        # Check if all in index are of the same type
        if not all(isinstance(i, index_type) for i in index):
            raise ValueError(f'{message}')

        # Check type and converts to DATETIME_FORMAT_WITH_TZ
        if index_type is str:
            # Validate format of string and return error if not valid
            try:
                pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
            except ValueError as e:
                raise ValueError(f'{message}') from e

        elif index_type == datetime or index_type == pd.Timestamp:
            data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)  # type: ignore[attr-defined]
        else:
            raise ValueError(f'{message}')

        return data

    def transform(
        self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
    ) -> dict:
        """
        Transform data using a model.

        Parameters:
            - model_name (str): The name of the model to use for transformation.
            - data (pandas.DataFrame): The data to transform.
            - model_retention (int): The number of minutes to keep the model.

        Returns:
            - dict: A dictionary containing the transformed data.
        """

        try:
            self.logger.custom_debug(
                f'Data received for model transformation: {data.to_csv()}', metadata
            )

            model_retention = model_config.get('retention_minutes', 0)
            flavor = model_config.get('transform_flavor', 'sklearn')
            compressed = model_config.get('is_compressed', False)
            retention_target = model_config.get('retention_target', 'model')
            transform_keyword = model_config.get('transform_function_keyword', 'predict')

            transformed_data = self.model_serving.get_cached_transform(
                model_name,
                data,
                model_retention,
                flavor,
                compressed,
                retention_target,
                transform_keyword,
            )

            self.logger.custom_debug(
                f'Data received from model transformation: {transformed_data.to_csv()}', metadata
            )

            transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata)

            return {'success': True, 'content': transformed_data.to_dict()}

        except Exception as e:  # noqa: BLE001
            return {
                'success': False,
                'content': {'message': str(e), 'traceback': traceback.format_exc()},
            }

    def predict(
        self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
    ) -> dict:
        """
        Predict data using a model.

        Parameters:
            - model_name (str): The name of the model to use for prediction.
            - data (pandas.DataFrame): The data to predict.
            - model_retention (int): The number of minutes to keep the model.

        Returns:
            - dict: A dictionary containing the predicted data.
        """
        try:
            model_retention = model_config.get('retention_minutes', 0)
            flavor = model_config.get('predict_flavor', 'pyfunc')
            compressed = model_config.get('is_compressed', False)
            retention_target = model_config.get('retention_target', 'model')

            input_index = data.index
            start_time = datetime.now()

            self.logger.custom_debug(
                f'Data received for model prediction: {data.to_csv()}', metadata
            )
            data = self.model_serving.get_cached_predict(
                model_name, data, model_retention, flavor, compressed, retention_target
            )

            end_time = datetime.now()
            data = pd.DataFrame(data, columns=['prediction'])
            self.logger.custom_debug(
                f'Data received from model prediction: {data.to_csv()}', metadata
            )
            data.index = input_index
            data['response_time'] = (end_time - start_time).total_seconds()

            return {'success': True, 'content': data.to_dict()}

        except Exception as e:  # noqa: BLE001
            return {
                'success': False,
                'content': {'message': str(e), 'traceback': traceback.format_exc()},
            }

    def get_experiment_by_run_id(self, run_id: str) -> dict:
        # Get the run information using the run_id
        run = mlflow.get_run(run_id)

        # Extract the experiment ID from the run
        experiment_id = run.info.experiment_id

        # Get the experiment details using the experiment ID
        experiment = mlflow.get_experiment(experiment_id)
        experiment_name = experiment.name
        return experiment_name

    def get_next_run_name(self, model_name: str) -> str:
        """
        Generate the next run name for a specific MLFlow model.

        This method calculates the next sequential run number for a model
        by searching existing runs and incrementing the count. It ensures
        unique run names for model training and retraining operations.

        Args:
            model_name (str): The name of the MLFlow model

        Returns:
            str: The next run name in format 'model_name-run_number'
        """
        runs = mlflow.search_runs(experiment_names=[model_name], order_by=['start_time desc'])
        next_run_number = len(runs) + 1
        return f'{model_name}-{next_run_number}'

    def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
        """
        Create a new MLFlow experiment for model retraining.

        This method sets up the complete environment for model retraining by:
        1. Loading the current production prediction model
        2. Loading the current production transformation model
        3. Fitting the transformation model with new data
        4. Preparing data for prediction model retraining
        5. Setting up the MLFlow experiment context

        Args:
            model_name (str): Name of the MLFlow model to retrain
            data (pd.DataFrame): Training data for model retraining

        Returns:
            tuple: (prediction_model, data_model, experiment)
                - prediction_model: Loaded prediction model for retraining
                - data_model: Fitted transformation model
                - experiment: MLFlow experiment name
        """
        # load predictor model
        predictor_uri = f'models:/{model_name}/production'
        # load transform model
        latest_production_id = self.model_serving.get_model_info(model_name)  # type: ignore[no-any-return]
        transform_uri = self.model_serving.get_model_uri(latest_production_id, prediction=False)
        # load
        data_model = mlflow.sklearn.load_model(transform_uri)
        prediction_model = mlflow.sklearn.load_model(predictor_uri)
        data_model = data_model.fit(data)
        treated_data = data_model.predict(data)

        target_name = data_model.target_variable
        y = data[target_name]
        treated_data = pd.merge(treated_data, y, left_index=True, right_index=True)
        prediction_model = prediction_model.fit(treated_data)
        experiment = self.get_experiment_by_run_id(latest_production_id)
        mlflow.set_experiment(experiment)

        return prediction_model, data_model, experiment

    def perform_model_retrain(
        self, prediction_model, data_model, experiment: str, model_name: str, data: pd.DataFrame
    ):
        """
        Execute the complete model retraining process in MLFlow.

        This method performs the actual model retraining by:
        1. Starting a new MLFlow run with descriptive metadata
        2. Logging model parameters and hyperparameters
        3. Retraining both prediction and transformation models
        4. Logging training data as artifacts
        5. Saving retrained models to MLFlow registry

        Args:
            prediction_model: MLFlow prediction model to retrain
            data_model: MLFlow transformation model to retrain
            experiment (str): MLFlow experiment name for the retraining
            model_name (str): Name of the model being retrained
            data (pd.DataFrame): Training data used for retraining

        Returns:
            tuple: (status_message, experiment_name)
                - status_message (str): Success confirmation message
                - experiment_name (str): Name of the experiment
        """
        pred_model_atributes = vars(prediction_model)  # load class attributes
        data_model_atributes = vars(data_model)  # load class attributes
        experiment_description = f'Retrain model {model_name} with new data'
        current_run_name = self.get_next_run_name(experiment)
        with mlflow.start_run(
            run_name=current_run_name, description=experiment_description
        ) as _run:
            # update transfomation model
            # fixed parameters
            for name_atribute, val_atribute in pred_model_atributes.items():
                if name_atribute != 'model':
                    mlflow.log_param(name_atribute, val_atribute)
            # update prediction model
            for name_atribute, val_atribute in data_model_atributes.items():
                if name_atribute != 'model':
                    mlflow.log_param(name_atribute, val_atribute)
            # dynamic parameters, including model itself
            mlflow.sklearn.log_model(data_model, 'data_model')

            makedirs('temp', exist_ok=True)

            file_path = f'temp/raw_data_{model_name}.csv'
            data.to_csv(file_path, index=True)

            # log the data raw
            mlflow.log_artifact(file_path)

            # dynamic parameters, including model itself
            mlflow.sklearn.log_model(prediction_model, 'prediction_model')
            mlflow.log_param('retrain', True)

            # clear temp file
            if path.exists(file_path):
                remove(file_path)

            return 'Model retrained successfully', experiment

    def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
        """
        Orchestrate the complete model retraining workflow.

        This method coordinates the entire model retraining process by:
        1. Creating the MLFlow experiment environment
        2. Loading existing production models
        3. Executing the retraining process
        4. Returning comprehensive retraining results

        Args:
            data (pd.DataFrame): Training data for model retraining
            model_name (str): Name of the MLFlow model to retrain

        Returns:
            tuple: (status_message, experiment_name)
                - status_message (str): Retraining operation status
                - experiment_name (str): MLFlow experiment identifier
        """
        prediction_model, data_model, experiment = self.create_model_experiment(model_name, data)
        retrain_result = self.perform_model_retrain(
            prediction_model, data_model, experiment, model_name, data
        )
        return retrain_result

    def get_experiment(self, experiment_name: str) -> int:
        """
        Retrieve MLFlow experiment ID by experiment name.

        This method searches for an MLFlow experiment by name and
        returns its unique identifier. It provides error handling
        for non-existent experiments.

        Args:
            experiment_name (str): Name of the MLFlow experiment

        Returns:
            int: MLFlow experiment ID

        Raises:
            ValueError: If the experiment name is not found
        """
        experiment = mlflow.get_experiment_by_name(experiment_name)
        if experiment is None:
            raise ValueError(f'Experiment {experiment_name} not found')

        return experiment.experiment_id  # type: ignore[no-any-return]

    def get_experiment_last_run(self, experiment_id: int) -> str:
        """
        Retrieve the most recent retraining run ID for an experiment.

        This method searches for the latest run in an MLFlow experiment
        that has been marked as a retraining run. It filters runs by
        the 'retrain' parameter and orders them by completion time.

        Args:
            experiment_id (int): MLFlow experiment ID

        Returns:
            str: MLFlow run ID of the most recent retraining run

        Raises:
            ValueError: If runs data is not in expected DataFrame format
        """
        runs = mlflow.search_runs(
            experiment_ids=[experiment_id],
            filter_string='',  # Sem filtro no MLflow ainda
            output_format='pandas',
        )

        if not isinstance(runs, pd.DataFrame):
            raise ValueError('Runs is not a pandas DataFrame')

        # Filtrar apenas as runs onde params.retrain == True
        filtered_runs = runs[runs['params.retrain'] == 'True']

        # Converter a coluna 'end_time' para datetime
        filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])

        # Ordenar o DataFrame de forma descendente pela coluna 'end_time'
        filtered_runs = filtered_runs.sort_values(by='end_time', ascending=False)

        # Pegar a última run_id do DataFrame filtrado e ordenado
        latest_run_id = filtered_runs.iloc[0]['run_id']

        return latest_run_id

    def update_production_model_by_run_id(self, run_id: str, model_name: str) -> dict:
        """
        Update production model with a specific MLFlow run.

        This method promotes a model from a specific MLFlow run to
        production stage. It handles model registration, versioning,
        and stage transitions with proper error handling.

        Args:
            run_id (str): MLFlow run ID containing the model to promote
            model_name (str): Name of the MLFlow model

        Returns:
            dict: Model update metadata containing:
                - model_name (str): Name of the updated model
                - version (str): New model version number
                - mlflow_run_id (str): Source run ID

        Update Process:
            1. Registers the model from the specified run
            2. Retrieves the latest model version
            3. Transitions the model to 'Production' stage
            4. Archives existing production versions
        """
        # Registrar o modelo
        # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
        # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
        mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name)

        # Colocar a versão do modelo em produção
        # Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
        client = mlflow.tracking.MlflowClient()

        # Obter a versão mais recente registrada do modelo
        model_versions = client.get_registered_model(model_name).latest_versions

        if not isinstance(model_versions, list):
            raise ValueError('Model versions is not a list')

        max_version = max(model_versions, key=lambda x: int(x.version)).version

        # Mover a versão mais recente do modelo para o estágio de 'Production'
        client.transition_model_version_stage(
            name=model_name, version=max_version, stage='Production', archive_existing_versions=True
        )

        return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id}

    def update_production_model(self, experiment: str, model_name: str) -> dict:
        """
        Update production model using the latest retraining run.

        This method orchestrates the complete production model update
        process by identifying the most recent retraining run and
        promoting it to production stage.

        Args:
            experiment (str): MLFlow experiment name
            model_name (str): Name of the MLFlow model

        Returns:
            dict: Complete model update metadata containing:
                - model_name (str): Name of the updated model
                - version (str): New model version number
                - mlflow_run_id (str): Source run ID
                - mlflow_experiment_id (int): Experiment ID
        """
        experiment_id = self.get_experiment(experiment)
        run_id = self.get_experiment_last_run(experiment_id)
        metadata = self.update_production_model_by_run_id(run_id, model_name)

        metadata['mlflow_experiment_id'] = experiment_id

        return metadata

    def get_next_run_name_new(self, experiment_name: str) -> str:
        """
        Generates the next run name for a given experiment.

        Args:
            experiment_name (str): The name of the experiment for which the next run name is being generated.

        Returns:
            str: A unique run name in the format "<experiment_name>-<next_run_number>".
        """
        runs = self.model_serving.search_runs_by_name(
            experiment_names=[experiment_name], order_by=['start_time desc']
        )

        next_run_number = len(runs) + 1
        return f'{experiment_name}-{next_run_number}'

    def generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
        """
        Generates and organizes artifacts related to the training process, such as reports and data files.

        Args:
            data: The training model result containing the datasets, model, and parameters.

        Returns:
            The updated result object with paths to the generated artifacts.

        Raises:
            FileNotFoundError: If the reports directory or header.html file does not exist.
            ValueError: If run_name is not set.
        """
        # Validate that run_name is set
        if not data.run_name:
            error_msg = 'run_name must be set before generating artifacts'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        reference_data, current_data = self._init_artifacts_data(data)
        base_path = self._get_reports_directory()

        # Validate that reports directory exists
        if not path.exists(base_path):
            error_msg = f'Reports directory does not exist: {base_path}'
            self.logger.error(error_msg)
            raise FileNotFoundError(error_msg)

        data.run_dir = self._create_run_directory(base_path, data.run_name)
        header_file_path = path.join(base_path, 'header.html')

        # Validate that header.html exists
        if not path.exists(header_file_path):
            error_msg = f'Header file does not exist: {header_file_path}'
            self.logger.error(error_msg)
            raise FileNotFoundError(error_msg)

        self._setup_run_directory(data.run_dir, header_file_path)
        return self._generate_report(reference_data, current_data, data)

    def save_run(self, data: TrainModelResult):
        """
        Logs the details of a machine learning run, including parameters, metrics, models, and artifacts,
        to the Sientia tracking system.

        Args:
            data: The training model result containing the datasets, model, parameters,
                  and evaluation metrics.

        Raises:
            ValueError: If required metrics or artifacts are missing.
            Exception: If MLflow logging fails for any reason.
        """
        # Validate that required artifacts exist before attempting to log
        if not data.report_path or not path.exists(data.report_path):
            error_msg = f'Report file does not exist: {data.report_path}'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if not data.train_data_path or not path.exists(data.train_data_path):
            error_msg = f'Training data file does not exist: {data.train_data_path}'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if not data.test_data_path or not path.exists(data.test_data_path):
            error_msg = f'Test data file does not exist: {data.test_data_path}'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        # Validate that metrics are present
        if data.mse_val is None or data.r2_val is None or data.mae_val is None:
            error_msg = 'One or more metrics (MSE, R2, MAE) are None'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        try:
            # Prepare parameters
            train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
            interval_strs = [
                (str(interval[0]), str(interval[1]))
                for interval in (data.params.removed_intervals or [])
            ]

            # Set experiment and create run
            self.model_serving.set_experiment(data.params.experiment_name)
            self.logger.info(
                f"Logging run '{data.run_name}' to experiment '{data.params.experiment_name}'"
            )

            with self.model_serving.save_experiment(
                run_name=data.run_name, description=data.params.experiment_name
            ):
                # Log model parameters
                self.model_serving.log_param('model_type', 'Linear Regression')
                self.model_serving.log_param('target_variable', data.params.target_variable)
                self.model_serving.log_param('input_variables', data.params.variable_columns)
                self.model_serving.log_param('lag_train', data.params.lag_train)
                self.model_serving.log_param('lag_val', data.params.lag_val)
                self.model_serving.log_param('ma', data.params.window)
                self.model_serving.log_param('low_lim', data.params.low_lim)
                self.model_serving.log_param('upp_lim', data.params.upp_lim)
                self.model_serving.log_param('normalized', data.scaler_dict)
                self.model_serving.log_param('ar', data.params.include_ar)
                self.model_serving.log_param('Train_test_split', train_test_split)
                self.model_serving.log_param('Removed_intervals', interval_strs)
                self.model_serving.log_param('Retrain', False)

                # Log evaluation metrics
                self.model_serving.log_metric('MSE', data.mse_val)
                self.model_serving.log_metric('R2', data.r2_val)
                self.model_serving.log_metric('MAE', data.mae_val)

                # Log models
                self.model_serving.log_model(data.process_data, 'data_model')
                self.model_serving.log_model(data.regr, 'prediction_model')

                # Log artifacts
                self.model_serving.log_artifact(data.report_path)
                self.model_serving.log_artifact(data.train_data_path)
                self.model_serving.log_artifact(data.test_data_path)

                self.logger.info(
                    f"Successfully logged run '{data.run_name}' with metrics: MSE={data.mse_val:.4f}, R2={data.r2_val:.4f}, MAE={data.mae_val:.4f}"
                )

        except Exception as e:
            error_msg = f"Failed to save run '{data.run_name}' to MLflow: {str(e)}"
            self.logger.error(error_msg)
            raise Exception(error_msg) from e

    def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
        """
        Prepares the reference and current datasets for artifact generation.

        Args:
            data: The training model result containing the datasets and model.

        Returns:
            tuple: A tuple containing:
                - reference_data: The training dataset with predictions added.
                - current_data: The testing dataset with predictions added.

        Raises:
            ValueError: If training or test datasets are empty or invalid.
            AttributeError: If required attributes are missing from the data object.
        """
        # Validate that required DataFrames are not empty
        # Note: x_train, y_train, x_test, y_test, and regr are required fields in TrainModelResult
        # so we only check if they are empty, not None
        if data.x_train.empty:
            error_msg = 'Training features (x_train) are empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if data.y_train.empty:
            error_msg = 'Training target (y_train) is empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if data.x_test.empty:
            error_msg = 'Test features (x_test) are empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        if data.y_test.empty:
            error_msg = 'Test target (y_test) is empty'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        # Validate that predictions exist (y_pred is optional, so check for None)
        if data.y_pred is None:
            error_msg = 'Test predictions (y_pred) are None'
            self.logger.error(error_msg)
            raise ValueError(error_msg)

        # Prepare reference data (training set)
        reference_data = pd.concat([data.x_train, data.y_train], axis=1)
        reference_data = reference_data.rename(columns={data.params.target_variable: 'target'})
        reference_data['prediction'] = data.regr.predict(data.x_train)

        # Prepare current data (test set)
        current_data = pd.concat([data.x_test, data.y_test], axis=1)
        current_data = current_data.rename(columns={data.params.target_variable: 'target'})
        current_data['prediction'] = data.y_pred

        return reference_data, current_data

    def _create_run_directory(self, base_path: str, run_name: str) -> str:
        """
        Creates a directory inside the 'reports' folder with the run name and a timestamp.

        Uses microsecond precision in timestamp to minimize collision probability
        in high-concurrency scenarios.

        Args:
            base_path (str): The path to the 'reports' folder.
            run_name (str): The name of the run.

        Returns:
            str: The path to the created directory.

        Raises:
            PermissionError: If there are insufficient permissions to create the directory.
            OSError: If directory creation fails for any other reason.
        """
        # Use microsecond precision to reduce collision probability
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
        run_dir = path.join(base_path, f'{run_name}_{timestamp}')

        try:
            makedirs(run_dir, exist_ok=True)
            self.logger.info(f'Created run directory: {run_dir}')
            return run_dir
        except PermissionError as e:
            error_msg = f'Permission denied when creating directory: {run_dir}'
            self.logger.error(error_msg)
            raise PermissionError(error_msg) from e
        except OSError as e:
            error_msg = f'Failed to create directory {run_dir}: {str(e)}'
            self.logger.error(error_msg)
            raise OSError(error_msg) from e

    def _setup_run_directory(self, run_dir: str, header_file_path: str):
        """
        Creates empty files and copies a header file into the specified run directory.

        Note: Lock removed as each run has its own unique directory, so no synchronization
        is needed between different runs. File operations within the same directory are
        atomic at the OS level.

        Args:
            run_dir (str): The path to the run directory where the files will be created.
            header_file_path (str): The path to the header.html file to be copied.

        Raises:
            FileNotFoundError: If the header file does not exist.
            PermissionError: If there are insufficient permissions to create files.
            OSError: If file creation or copying fails for any other reason.
        """
        empty_files = ['data_drift.html', 'data_quality.html', 'regression.html']

        try:
            # Create empty placeholder files
            for file_name in empty_files:
                file_path = path.join(run_dir, file_name)
                with open(file_path, 'w'):
                    pass  # Create empty file

            # Copy header file to run directory
            header_dest = path.join(run_dir, 'header.html')
            shutil.copy(header_file_path, header_dest)

            self.logger.info(f'Run directory setup completed successfully in: {run_dir}')

        except FileNotFoundError as e:
            error_msg = f'Header file not found: {header_file_path}'
            self.logger.error(error_msg)
            raise FileNotFoundError(error_msg) from e
        except PermissionError as e:
            error_msg = f'Permission denied when setting up directory: {run_dir}'
            self.logger.error(error_msg)
            raise PermissionError(error_msg) from e
        except OSError as e:
            error_msg = f'Failed to setup run directory {run_dir}: {str(e)}'
            self.logger.error(error_msg)
            raise OSError(error_msg) from e

    def _generate_report(
        self, reference_data: pd.DataFrame, current_data: pd.DataFrame, data: TrainModelResult
    ) -> TrainModelResult:
        """
        Generates a comprehensive report summarizing data quality, data drift, and regression analysis.

        Args:
            reference_data (pd.DataFrame): The training dataset with predictions added.
            current_data (pd.DataFrame): The testing dataset with predictions added.
            data: The training model result containing the datasets, model, and parameters.

        Returns:
            The updated result object with paths to the generated report and data files.

        Raises:
            ValueError: If data conversion to float64 fails or DataFrames are invalid.
            PermissionError: If there are insufficient permissions to write files.
            OSError: If file writing fails for any other reason.
        """
        try:
            # Convert data to float64 for report generation
            # This may raise ValueError if data contains non-numeric values
            reference_data_float = reference_data.astype(np.float64)
            current_data_float = current_data.astype(np.float64)

            # Initialize report generator
            report = Reports(
                reference_data=reference_data_float,
                current_data=current_data_float,
                base_path=data.run_dir,
            )

            # Generate report sections
            report.add_data_quality_section(columns=data.params.variable_columns + ['target'])
            report.add_data_drift_section(columns=data.params.variable_columns + ['target'])
            report.add_regression_section()

            # Validate that run_dir is set (should be set by _create_run_directory)
            if not data.run_dir:
                error_msg = 'run_dir is not set after directory creation'
                self.logger.error(error_msg)
                raise ValueError(error_msg)

            # Save HTML report
            data.report_path = path.join(data.run_dir, 'report.html')
            report.save_all_sections_html(data.report_path)
            self.logger.info(f'Generated HTML report: {data.report_path}')

            # Save training data CSV
            data.train_data_path = path.join(data.run_dir, 'train_data.csv')
            reference_data.to_csv(data.train_data_path, index=False)
            self.logger.info(f'Saved training data: {data.train_data_path}')

            # Save test data CSV
            data.test_data_path = path.join(data.run_dir, 'test_data.csv')
            current_data.to_csv(data.test_data_path, index=False)
            self.logger.info(f'Saved test data: {data.test_data_path}')

            return data

        except ValueError as e:
            error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'
            self.logger.error(error_msg)
            raise ValueError(error_msg) from e
        except PermissionError as e:
            error_msg = f'Permission denied when writing report files to: {data.run_dir}'
            self.logger.error(error_msg)
            raise PermissionError(error_msg) from e
        except OSError as e:
            error_msg = f'Failed to generate report in {data.run_dir}: {str(e)}'
            self.logger.error(error_msg)
            raise OSError(error_msg) from e

    def _get_reports_directory(self) -> str:
        """
        Get the absolute path to the reports directory.

        Returns:
            str: Absolute path to model_manager/reports directory.
        """
        # Get the directory where this file is located (model_manager/utils/repository/)
        current_file_dir = path.dirname(path.abspath(__file__))
        # Navigate up to model_manager/ and then to reports/
        model_manager_dir = path.dirname(path.dirname(current_file_dir))
        reports_dir = path.join(model_manager_dir, 'reports')
        return reports_dir