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

"""
Training repository for ML model training operations.

This module provides the core training logic for machine learning models,
including data preprocessing, model training, and post-training calculations.
Migrated from laborious/utils/train_model_utils.py.
"""

from io import BytesIO

import numpy as np
import pandas as pd
from sientia_do.observability.logger import Logger
from sientia_do.operations.df_preprocessor import load_data
from sientia_do.operations.normalization import MinMaxScaler, Z_Scaler

from model_manager.sientia.metrics import mae, mse, r2
from model_manager.sientia.models import DataPreprocessor, LinearRegressionModel
from model_manager.sientia.utils import split_train_test
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult


class TrainingRepository:
    """
    Repository for machine learning model training operations.

    This class encapsulates the core logic for training ML models, migrated from
    laborious/utils/train_model_utils.py. Follows the same pattern as MLFlowRepository
    with instance methods and logger integration.

    Attributes:
        logger (Logger): Logger instance for observability and debugging
    """

    def __init__(self, logger: Logger):
        """
        Initialize TrainingRepository with logger.

        Args:
            logger: Logger instance for observability
        """
        self.logger = logger

    def train(self, uploaded_file: BytesIO, params: TrainModelParams) -> TrainModelResult:
        """
        Train a machine learning model using the provided file and parameters.

        This method orchestrates the training pipeline:
        1. Load data from BytesIO file
        2. Initialize and fit data preprocessor
        3. Transform data and validate
        4. Split into train/test sets
        5. Initialize scaler dictionary
        6. Train LinearRegression model

        Args:
            uploaded_file: BytesIO object containing training data (CSV format)
            params: Training parameters (TrainModelParams)

        Returns:
            TrainModelResult: Object containing trained model, processed data,
                            train/test splits, and scaler dictionary

        Raises:
            ValueError: If transformed data is empty
            Exception: If data loading, preprocessing, or training fails
        """
        data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
        process_data = self._init_data_preprocessor(params)
        process_data.fit(data)
        data_view = process_data.transform(data)

        if len(data_view) <= 0:
            raise ValueError('Data view is empty after transformation')

        x_train, x_test, y_train, y_test = split_train_test(
            data_view[params.variable_columns],
            data_view[params.target_variable],
            train_size=params.train_size / 100,
            shuffle=params.shuffle,
            random_state=42,
        )

        data_train = pd.concat([x_train, y_train], axis=1)
        scaler_dict = self._init_scaler_dict(process_data, params)

        regr = LinearRegressionModel(
            target_variable=params.target_variable,
            variable_columns=params.variable_columns,
            degree=params.degree,
            interaction_only=params.interaction_only,
        )

        regr.fit(data_train)
        self.logger.info(
            f'Model trained successfully - experiment run id: {params.experiment_run_id}'
        )

        return TrainModelResult(
            params=params,
            process_data=process_data,
            x_train=x_train,
            x_test=x_test,
            y_train=y_train,
            y_test=y_test,
            regr=regr,
            scaler_dict=scaler_dict,
        )

    def after_train_calculation(
        self, params: TrainModelParams, tmr: TrainModelResult
    ) -> TrainModelResult:
        """
        Perform post-training calculations: predictions, denormalization, and metrics.

        This method completes the training pipeline by:
        1. Making predictions on test set
        2. Denormalizing all data (if scaler was used)
        3. Reordering data by index
        4. Calculating evaluation metrics (MSE, MAE, R²)

        Args:
            params: Training parameters used during model training
            tmr: Result object from training

        Returns:
            TrainModelResult: Updated result with predictions, denormalized data,
                            and metrics (mse_val, mae_val, r2_val)
        """
        y_pred_array = tmr.regr.predict(tmr.x_test)

        if params.use_scaler:
            scaler = tmr.process_data.get_scaler()

            # If using custom scaler with denormalize_* helpers
            if hasattr(scaler, 'denormalize_single_input'):
                for col in params.variable_columns:
                    tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
                    tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)

                tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
                tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
                y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
            else:
                # Fallback for sklearn StandardScaler: only inverse-transform features
                feature_cols = getattr(
                    tmr.process_data, 'feature_names_order', params.variable_columns
                )
                # Ensure columns are in the same order used during fit
                x_train_features = tmr.x_train[feature_cols]
                x_test_features = tmr.x_test[feature_cols]

                tmr.x_train[feature_cols] = scaler.inverse_transform(x_train_features)
                tmr.x_test[feature_cols] = scaler.inverse_transform(x_test_features)
                # Target was not scaled with StandardScaler in preprocessing; leave y as-is

        tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
        tmr.y_pred.name = f'{params.target_variable}_pred'

        tmr.x_train = tmr.x_train.sort_index()
        tmr.x_test = tmr.x_test.sort_index()
        tmr.y_train = tmr.y_train.sort_index()
        tmr.y_test = tmr.y_test.sort_index()
        tmr.y_pred = tmr.y_pred.sort_index()

        assert tmr.y_pred is not None, 'y_pred should be set at this point'

        tmr.mse_val = round(
            mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
            2,
        )

        tmr.mae_val = round(
            mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
            2,
        )

        tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)

        # Extract model equation
        tmr.equation = self._extract_model_equation(tmr.regr, params)

        self.logger.info(
            f'Model metrics calculated successfully - experiment run id: {params.experiment_run_id}'
        )
        return tmr

    def _init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
        """
        Initialize dictionary containing scaling parameters for features and target.

        This method extracts scaling parameters from the fitted scaler to enable
        denormalization of predictions and debugging of the normalization process.

        Args:
            process_data: Fitted DataPreprocessor object with scaler
            params: Training parameters including scaler configuration

        Returns:
            dict: Scaling parameters for each feature and target variable.
                 Structure depends on scaler type:
                 - MinMaxScaler: {'feature': {'min': float, 'max': float}, ...}
                 - Z_Scaler: Dictionary from scaler.create_dict()
                 - Empty dict: If no scaler is used

        Raises:
            AttributeError: If scaler doesn't have expected attributes
        """
        scaler_dict = {}

        if params.use_scaler:
            scaler = process_data.get_scaler()

            if isinstance(scaler, MinMaxScaler):
                # Extract min/max for each feature
                for i, col in enumerate(params.variable_columns):
                    scaler_dict[col] = {'min': scaler.x_min[i], 'max': scaler.x_max[i]}

                # Extract min/max for target variable
                scaler_dict[params.target_variable] = {
                    'min': scaler.y_min,
                    'max': scaler.y_max,
                }

            elif isinstance(scaler, Z_Scaler):
                scaler_dict = scaler.create_dict()

        return scaler_dict

    def _init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
        """
        Initialize DataPreprocessor with training parameters.

        Args:
            params: Training parameters containing preprocessor configuration

        Returns:
            DataPreprocessor: Configured preprocessor ready for fitting
        """
        # Convert removed_intervals to list of tuples if needed
        removed_intervals = None
        if params.removed_intervals:
            removed_intervals = [
                (interval[0], interval[1]) if isinstance(interval, (list, tuple)) else interval
                for interval in params.removed_intervals
            ]

        return DataPreprocessor(
            target_variable=params.target_variable,
            input_columns=params.variable_columns,
            nan_treatment=params.nan_treatment,
            lag_train=params.lag_train,
            lag_transform=params.lag_val,
            start_date=params.start_date,
            end_date=params.end_date,
            removed_intervals=removed_intervals,
            static_threshold=1 if params.rem_static_win else None,
            low_lim=params.low_lim,
            upp_lim=params.upp_lim,
            scaler_name=params.scaler_name,
            scaler_params={} if params.use_scaler else None,
            ar_var=params.target_variable if params.include_ar else None,
        )

    def _extract_model_equation(
        self, regr: LinearRegressionModel, params: TrainModelParams
    ) -> dict:
        """
        Extract the linear regression equation coefficients and create equation metadata.

        This method extracts the coefficients and intercept from the trained model
        and creates a structured dictionary containing the equation information
        for serialization as JSON artifact.

        Args:
            regr: Trained LinearRegressionModel object
            params: Training parameters containing variable information

        Returns:
            dict: Equation metadata containing:
                - target_variable: Name of the target variable
                - coefficients: Dictionary mapping variable names to coefficients
                - intercept: Model intercept value
                - equation_string: Human-readable equation string
                - latex_equation: LaTeX formatted equation
        """
        coefficients = regr.regr.coef_
        intercept = regr.regr.intercept_

        # Get feature names - for polynomial models, use poly_feature_names
        if params.degree > 1 and regr.poly_feature_names:
            feature_names = regr.poly_feature_names
        else:
            feature_names = params.variable_columns

        # Create coefficients dictionary
        coefficients_dict = {}
        for i, var in enumerate(feature_names):
            if i < len(coefficients):
                coefficients_dict[var] = float(coefficients[i])

        # Create equation string
        equation_parts = [f'{coef:.6f} * {var}' for var, coef in coefficients_dict.items()]
        equation_string = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(
            equation_parts
        )

        # Create LaTeX equation
        latex_parts = [f'{coef:.6f} \\cdot {var}' for var, coef in coefficients_dict.items()]
        latex_equation = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(latex_parts)

        return {
            'target_variable': params.target_variable,
            'coefficients': coefficients_dict,
            'intercept': float(intercept),
            'equation_string': equation_string,
            'latex_equation': latex_equation,
            'model_type': params.model_name,
            'degree': params.degree,
            'interaction_only': params.interaction_only,
            'original_features': params.variable_columns,
        }