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


def _frontend_format_to_strftime(fmt: str) -> str:
    """Convert front-end date format (dd/MM/yyyy HH:mm:ss) to Python strftime (%d/%m/%Y %H:%M:%S)."""
    if not fmt:
        return fmt
    out = fmt.replace('yyyy', '%Y').replace('MM', '%m').replace('dd', '%d')
    out = out.replace('HH', '%H').replace('mm', '%M').replace('ss', '%S')
    return out


def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame:
    """If date_column and date_format are set, parse the column as datetime to avoid comparison errors downstream."""
    if not params.date_column or not params.date_format or params.date_column not in data.columns:
        return data
    try:
        python_fmt = _frontend_format_to_strftime(params.date_format)
        data = data.copy()
        data[params.date_column] = pd.to_datetime(
            data[params.date_column], format=python_fmt, errors='coerce'
        )
    except Exception as e:
        raise ValueError(
            f'Failed to parse date column "{params.date_column}" with format "{params.date_format}": {e}'
        ) from e
    return data


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)
        if data is None:
            raise ValueError(
                'Failed to load CSV data: load_data returned None. '
                'Check file encoding, line separator and decimal separator.'
            )
        data = _ensure_date_column_parsed(data, params)
        data = self._configure_datetime_index(data, params)

        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 _get_static_threshold(self, params: TrainModelParams) -> int | None:
        """
        Get the static threshold value based on parameters.

        Args:
            params: Training parameters containing static window configuration

        Returns:
            int | None: Static threshold value (1-1000) if rem_static_win is True, None otherwise
        """
        if not params.rem_static_win:
            return None
        return params.static_threshold if params.static_threshold is not None else 1

    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,
            date_format=params.date_format,
            removed_intervals=removed_intervals,
            static_threshold=self._get_static_threshold(params),
            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 | None, params: TrainModelParams
    ) -> pd.DataFrame:
        """
        Configure datetime index for the DataFrame.

        Guards against None to avoid 'NoneType' object has no attribute 'index' downstream.
        Prefers params.date_column when set; otherwise looks for common timestamp column names.
        """
        if data is None:
            raise ValueError(
                'Data is None after load_data. '
                'Check file format, line separator and decimal separator.'
            )
        if not isinstance(data, pd.DataFrame):
            raise TypeError(f'Expected DataFrame, got {type(data).__name__}')

        if isinstance(data.index, pd.DatetimeIndex):
            self.logger.info('DataFrame already has DatetimeIndex')
            return data.sort_index()

        common_timestamp_columns = [
            'timestamp',
            'Timestamp',
            'TIMESTAMP',
            'date',
            'Date',
            'DATE',
            'DATA',
            'datetime',
            'DateTime',
        ]
        timestamp_columns = (
            [params.date_column] if params.date_column else []
        ) + [c for c in common_timestamp_columns if c != params.date_column]

        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