"""
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)

        # Configure datetime index if timestamp column exists
        # Required for TimeSeriesDiscontinuityAnalyzer (static window removal)
        data = self._configure_datetime_index(data)

        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)
        """
        # Calculate predictions BEFORE denormalization (important for polynomial models)
        y_pred_array = tmr.regr.predict(tmr.x_test)
        y_train_pred_array = tmr.regr.predict(tmr.x_train)

        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)
                y_train_pred_array = scaler.denormalize_predictions(
                    y_train_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.y_train_pred = pd.Series(y_train_pred_array, index=tmr.y_train.index)
        tmr.y_train_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()
        tmr.y_train_pred = tmr.y_train_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=(params.static_threshold if params.static_threshold is not None else 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,
        }

    def _configure_datetime_index(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        Configure datetime index for the DataFrame.

        This method attempts to identify a timestamp column and set it as the
        DataFrame index with DatetimeIndex type. This is required for
        TimeSeriesDiscontinuityAnalyzer (used in static window removal).

        The method looks for common timestamp column names and converts the
        first matching column to datetime, then sets it as the index.

        Args:
            data: Input DataFrame

        Returns:
            pd.DataFrame: DataFrame with DatetimeIndex if timestamp column found,
                         otherwise returns original DataFrame unchanged
        """
        # If index is already DatetimeIndex, just ensure it's sorted
        if isinstance(data.index, pd.DatetimeIndex):
            self.logger.info('DataFrame already has DatetimeIndex')
            return data.sort_index()

        # Common timestamp column names
        timestamp_columns = [
            'timestamp',
            'Timestamp',
            'TIMESTAMP',
            'date',
            'Date',
            'DATE',
            'datetime',
            'DateTime',
        ]

        for col in timestamp_columns:
            if col in data.columns:
                try:
                    data[col] = pd.to_datetime(data[col])
                    data = data.set_index(col)
                    data = data.sort_index()
                    self.logger.info(f'Configured datetime index from column: {col}')
                    return data
                except (ValueError, TypeError) as e:
                    self.logger.warning(f'Failed to convert column {col} to datetime: {e}')
                    continue

        # If no timestamp column found, check if first column looks like a timestamp
        first_col = data.columns[0]
        try:
            # Try to parse first column as datetime
            test_values = data[first_col].head(10).dropna()
            if len(test_values) > 0:
                pd.to_datetime(test_values)
                data[first_col] = pd.to_datetime(data[first_col])
                data = data.set_index(first_col)
                data = data.sort_index()
                self.logger.info(f'Configured datetime index from first column: {first_col}')
                return data
        except (ValueError, TypeError):
            pass

        self.logger.warning('No timestamp column found - some features may not work correctly')
        return data