"""Unit tests for TrainingRepository."""

from io import BytesIO
from unittest.mock import MagicMock, patch

import numpy as np
import pandas as pd
import pytest

from model_manager.sientia.models import LinearRegressionModel
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
from model_manager.utils.repository.training_repository import TrainingRepository


@pytest.fixture
def mock_logger():
    """Create a mock logger for testing."""
    return MagicMock()


@pytest.fixture
def training_repo(mock_logger):
    """Create TrainingRepository instance with mock logger."""
    return TrainingRepository(logger=mock_logger)


@pytest.fixture
def sample_params():
    """Create sample TrainModelParams for testing."""
    return TrainModelParams(
        experiment_run_id=1,
        experiment_name='test_experiment',
        target_variable='target',
        variable_columns=['var1', 'var2', 'var3'],
        lag_train=0,
        lag_val=0,
        rem_static_win=False,
        low_lim={},
        upp_lim={},
        window=0,
        use_scaler=False,
        include_ar=False,
        train_size=80,
        shuffle=True,
        bucket_name='test-bucket',
        file_name='test.csv',
        line_separator=',',
        decimal_separator='.',
        removed_intervals=[],
    )


@pytest.fixture
def sample_linear_model():
    """Create a mock LinearRegressionModel with known coefficients."""
    mock_model = MagicMock(spec=LinearRegressionModel)
    mock_sklearn_model = MagicMock()
    mock_sklearn_model.coef_ = np.array([1.5, -0.75, 2.25])
    mock_sklearn_model.intercept_ = 10.5
    mock_model.regr = mock_sklearn_model
    return mock_model


class TestExtractModelEquation:
    """Tests for _extract_model_equation method."""

    def test_extract_equation_basic(self, training_repo, sample_params, sample_linear_model):
        """Test basic equation extraction with simple coefficients."""
        result = training_repo._extract_model_equation(sample_linear_model, sample_params)

        assert result['target_variable'] == 'target'
        assert result['model_type'] == 'Linear Regression'
        assert result['intercept'] == 10.5

        assert 'var1' in result['coefficients']
        assert 'var2' in result['coefficients']
        assert 'var3' in result['coefficients']

        assert result['coefficients']['var1'] == 1.5
        assert result['coefficients']['var2'] == -0.75
        assert result['coefficients']['var3'] == 2.25

    def test_extract_equation_string_format(
        self, training_repo, sample_params, sample_linear_model
    ):
        """Test equation string is formatted correctly."""
        result = training_repo._extract_model_equation(sample_linear_model, sample_params)

        expected_string = (
            'target = 10.500000 + 1.500000 * var1 + -0.750000 * var2 + 2.250000 * var3'
        )
        assert result['equation_string'] == expected_string

    def test_extract_equation_latex_format(self, training_repo, sample_params, sample_linear_model):
        """Test LaTeX equation is formatted correctly."""
        result = training_repo._extract_model_equation(sample_linear_model, sample_params)

        expected_latex = 'target = 10.500000 + 1.500000 \\cdot var1 + -0.750000 \\cdot var2 + 2.250000 \\cdot var3'
        assert result['latex_equation'] == expected_latex

    def test_extract_equation_single_variable(self, training_repo, sample_linear_model):
        """Test equation extraction with single variable."""
        params = TrainModelParams(
            experiment_run_id=1,
            experiment_name='test',
            target_variable='y',
            variable_columns=['x'],
            lag_train=0,
            lag_val=0,
            rem_static_win=False,
            low_lim={},
            upp_lim={},
            window=0,
            use_scaler=False,
            include_ar=False,
            train_size=80,
            shuffle=True,
            bucket_name='bucket',
            file_name='file.csv',
            line_separator=',',
            decimal_separator='.',
            removed_intervals=[],
        )

        # Mock model with single coefficient
        mock_model = MagicMock(spec=LinearRegressionModel)
        mock_sklearn_model = MagicMock()
        mock_sklearn_model.coef_ = np.array([3.14])
        mock_sklearn_model.intercept_ = 2.71
        mock_model.regr = mock_sklearn_model

        result = training_repo._extract_model_equation(mock_model, params)

        assert len(result['coefficients']) == 1
        assert result['coefficients']['x'] == 3.14
        assert result['intercept'] == 2.71
        assert 'y = 2.710000 + 3.140000 * x' == result['equation_string']

    def test_extract_equation_zero_coefficients(self, training_repo, sample_params):
        """Test equation extraction when coefficients are zero."""
        mock_model = MagicMock(spec=LinearRegressionModel)
        mock_sklearn_model = MagicMock()
        mock_sklearn_model.coef_ = np.array([0.0, 0.0, 0.0])
        mock_sklearn_model.intercept_ = 5.0
        mock_model.regr = mock_sklearn_model

        result = training_repo._extract_model_equation(mock_model, sample_params)

        assert all(v == 0.0 for v in result['coefficients'].values())
        assert result['intercept'] == 5.0

    def test_extract_equation_negative_intercept(self, training_repo, sample_params):
        """Test equation extraction with negative intercept."""
        mock_model = MagicMock(spec=LinearRegressionModel)
        mock_sklearn_model = MagicMock()
        mock_sklearn_model.coef_ = np.array([1.0, 2.0, 3.0])
        mock_sklearn_model.intercept_ = -5.5
        mock_model.regr = mock_sklearn_model

        result = training_repo._extract_model_equation(mock_model, sample_params)

        assert result['intercept'] == -5.5
        assert 'target = -5.500000 +' in result['equation_string']


class TestInitDataPreprocessor:
    """Tests for _init_data_preprocessor method."""

    def test_init_preprocessor_basic(self, training_repo, sample_params):
        """Test basic preprocessor initialization."""
        preprocessor = training_repo._init_data_preprocessor(sample_params)

        assert preprocessor.target_variable == 'target'
        assert preprocessor.input_columns == ['var1', 'var2', 'var3']

    def test_init_preprocessor_with_scaler(self, training_repo, sample_params):
        """Test preprocessor initialization with scaler enabled."""
        sample_params.use_scaler = True
        preprocessor = training_repo._init_data_preprocessor(sample_params)

        assert preprocessor.scaler_name == 'Standard Scaler'

    def test_init_preprocessor_without_scaler(self, training_repo, sample_params):
        """Test preprocessor initialization without scaler."""
        sample_params.use_scaler = False
        preprocessor = training_repo._init_data_preprocessor(sample_params)

        assert preprocessor.scaler_name == 'None'

    def test_init_preprocessor_with_ar(self, training_repo, sample_params):
        """Test preprocessor initialization with autoregressive variable."""
        sample_params.include_ar = True
        preprocessor = training_repo._init_data_preprocessor(sample_params)

        assert preprocessor.ar_var == 'target'

    def test_init_preprocessor_without_ar(self, training_repo, sample_params):
        """Test preprocessor initialization without autoregressive variable."""
        sample_params.include_ar = False
        preprocessor = training_repo._init_data_preprocessor(sample_params)

        assert preprocessor.ar_var is None

    def test_init_preprocessor_with_static_removal(self, training_repo, sample_params):
        """Test preprocessor with static window removal enabled."""
        sample_params.rem_static_win = True
        preprocessor = training_repo._init_data_preprocessor(sample_params)

        assert preprocessor.static_threshold == 1

    def test_init_preprocessor_lag_configuration(self, training_repo, sample_params):
        """Test preprocessor lag configuration."""
        sample_params.lag_train = 5
        sample_params.lag_val = 3
        preprocessor = training_repo._init_data_preprocessor(sample_params)

        # Check that lag dictionaries are created correctly
        for col in sample_params.variable_columns:
            assert preprocessor.lag_train[col] == 5
            assert preprocessor.lag_transform[col] == 3


class TestInitScalerDict:
    """Tests for _init_scaler_dict method."""

    def test_init_scaler_dict_without_scaler(self, training_repo, sample_params):
        """Test scaler dict initialization when scaler is not used."""
        from model_manager.sientia.models import DataPreprocessor

        sample_params.use_scaler = False
        process_data = MagicMock(spec=DataPreprocessor)

        result = training_repo._init_scaler_dict(process_data, sample_params)

        assert result == {}

    def test_init_scaler_dict_with_minmax_scaler(self, training_repo, sample_params):
        """Test scaler dict initialization with MinMaxScaler."""
        from sientia_do.operations.normalization import MinMaxScaler

        from model_manager.sientia.models import DataPreprocessor

        sample_params.use_scaler = True

        # Create mock MinMaxScaler
        mock_scaler = MagicMock(spec=MinMaxScaler)
        mock_scaler.x_min = np.array([0.0, 1.0, 2.0])
        mock_scaler.x_max = np.array([10.0, 11.0, 12.0])
        mock_scaler.y_min = 0.5
        mock_scaler.y_max = 100.5

        # Create mock preprocessor that returns the scaler
        process_data = MagicMock(spec=DataPreprocessor)
        process_data.get_scaler.return_value = mock_scaler

        result = training_repo._init_scaler_dict(process_data, sample_params)

        # Check feature scalers
        assert 'var1' in result
        assert 'var2' in result
        assert 'var3' in result
        assert result['var1'] == {'min': 0.0, 'max': 10.0}
        assert result['var2'] == {'min': 1.0, 'max': 11.0}
        assert result['var3'] == {'min': 2.0, 'max': 12.0}

        # Check target scaler
        assert 'target' in result
        assert result['target'] == {'min': 0.5, 'max': 100.5}

    def test_init_scaler_dict_with_z_scaler(self, training_repo, sample_params):
        """Test scaler dict initialization with Z_Scaler."""
        from sientia_do.operations.normalization import Z_Scaler

        from model_manager.sientia.models import DataPreprocessor

        sample_params.use_scaler = True

        # Create mock Z_Scaler
        mock_scaler = MagicMock(spec=Z_Scaler)
        expected_dict = {
            'var1': {'mean': 5.0, 'std': 1.5},
            'var2': {'mean': 10.0, 'std': 2.0},
            'target': {'mean': 50.0, 'std': 10.0},
        }
        mock_scaler.create_dict.return_value = expected_dict

        # Create mock preprocessor
        process_data = MagicMock(spec=DataPreprocessor)
        process_data.get_scaler.return_value = mock_scaler

        result = training_repo._init_scaler_dict(process_data, sample_params)

        assert result == expected_dict
        mock_scaler.create_dict.assert_called_once()


class TestAfterTrainCalculation:
    """Tests for after_train_calculation method."""

    @pytest.fixture
    def mock_train_result(self, sample_params, sample_linear_model):
        """Create a mock TrainModelResult."""
        x_train = pd.DataFrame(
            {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2]
        )
        x_test = pd.DataFrame({'var1': [10, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4])
        y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target')
        y_test = pd.Series([400, 500], index=[3, 4], name='target')

        # Mock predict to return a simple array
        sample_linear_model.predict = MagicMock(return_value=np.array([450.0, 550.0]))

        return TrainModelResult(
            params=sample_params,
            process_data=MagicMock(),
            x_train=x_train,
            x_test=x_test,
            y_train=y_train,
            y_test=y_test,
            regr=sample_linear_model,
            scaler_dict={},
        )

    def test_after_train_adds_predictions(self, training_repo, sample_params, mock_train_result):
        """Test that predictions are added to result."""
        result = training_repo.after_train_calculation(sample_params, mock_train_result)

        assert result.y_pred is not None
        assert len(result.y_pred) == len(result.y_test)
        assert result.y_pred.name == 'target_pred'

    def test_after_train_calculates_metrics(self, training_repo, sample_params, mock_train_result):
        """Test that metrics are calculated."""
        result = training_repo.after_train_calculation(sample_params, mock_train_result)

        assert result.mse_val is not None
        assert result.mae_val is not None
        assert result.r2_val is not None
        assert isinstance(result.mse_val, (int, float))
        assert isinstance(result.mae_val, (int, float))
        assert isinstance(result.r2_val, (int, float))

    def test_after_train_extracts_equation(self, training_repo, sample_params, mock_train_result):
        """Test that equation is extracted after training."""
        result = training_repo.after_train_calculation(sample_params, mock_train_result)

        assert result.equation is not None
        assert 'target_variable' in result.equation
        assert 'coefficients' in result.equation
        assert 'intercept' in result.equation
        assert 'equation_string' in result.equation
        assert 'latex_equation' in result.equation
        assert 'model_type' in result.equation

    def test_after_train_sorts_data(self, training_repo, sample_params, mock_train_result):
        """Test that data is sorted by index."""
        # Shuffle indices
        mock_train_result.x_train = mock_train_result.x_train.sample(frac=1)
        mock_train_result.y_train = mock_train_result.y_train.sample(frac=1)

        result = training_repo.after_train_calculation(sample_params, mock_train_result)

        assert result.x_train.index.is_monotonic_increasing
        assert result.y_train.index.is_monotonic_increasing
        assert result.x_test.index.is_monotonic_increasing
        assert result.y_test.index.is_monotonic_increasing

    def test_after_train_logs_success(
        self, training_repo, mock_logger, sample_params, mock_train_result
    ):
        """Test that success is logged."""
        training_repo.after_train_calculation(sample_params, mock_train_result)

        mock_logger.info.assert_called()
        assert any(
            'Model metrics calculated successfully' in str(call)
            for call in mock_logger.info.call_args_list
        )

    def test_after_train_with_custom_scaler_denormalization(
        self, training_repo, sample_params, sample_linear_model
    ):
        """Test denormalization with custom scaler that has denormalize methods."""
        sample_params.use_scaler = True

        # Create mock data
        x_train = pd.DataFrame(
            {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2]
        )
        x_test = pd.DataFrame({'var1': [10, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4])
        y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target')
        y_test = pd.Series([400, 500], index=[3, 4], name='target')

        # Mock predict to return a simple array
        sample_linear_model.predict = MagicMock(return_value=np.array([450.0, 550.0]))

        # Create mock scaler with denormalize methods
        mock_scaler = MagicMock()
        mock_scaler.denormalize_single_input = MagicMock(side_effect=lambda x, col: x * 2)
        mock_scaler.denormalize_predictions = MagicMock(return_value=np.array([900.0, 1100.0]))

        # Create mock preprocessor
        mock_process_data = MagicMock()
        mock_process_data.get_scaler.return_value = mock_scaler

        train_result = TrainModelResult(
            params=sample_params,
            process_data=mock_process_data,
            x_train=x_train,
            x_test=x_test,
            y_train=y_train,
            y_test=y_test,
            regr=sample_linear_model,
            scaler_dict={},
        )

        result = training_repo.after_train_calculation(sample_params, train_result)

        # Verify denormalize methods were called
        assert mock_scaler.denormalize_single_input.called
        assert mock_scaler.denormalize_predictions.called
        assert result.y_pred is not None

    def test_after_train_with_sklearn_scaler(
        self, training_repo, sample_params, sample_linear_model
    ):
        """Test denormalization with sklearn StandardScaler."""
        sample_params.use_scaler = True

        # Create mock data
        x_train = pd.DataFrame(
            {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2]
        )
        x_test = pd.DataFrame({'var1': [10, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4])
        y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target')
        y_test = pd.Series([400, 500], index=[3, 4], name='target')

        # Mock predict
        sample_linear_model.predict = MagicMock(return_value=np.array([450.0, 550.0]))

        # Create mock sklearn scaler (without denormalize methods)
        mock_scaler = MagicMock()
        # Remove denormalize methods to trigger sklearn path
        if hasattr(mock_scaler, 'denormalize_single_input'):
            delattr(mock_scaler, 'denormalize_single_input')
        mock_scaler.inverse_transform = MagicMock(side_effect=lambda x: x * 2)

        # Create mock preprocessor with feature_names_order
        mock_process_data = MagicMock()
        mock_process_data.get_scaler.return_value = mock_scaler
        mock_process_data.feature_names_order = ['var1', 'var2', 'var3']

        train_result = TrainModelResult(
            params=sample_params,
            process_data=mock_process_data,
            x_train=x_train,
            x_test=x_test,
            y_train=y_train,
            y_test=y_test,
            regr=sample_linear_model,
            scaler_dict={},
        )

        result = training_repo.after_train_calculation(sample_params, train_result)

        # Verify inverse_transform was called
        assert mock_scaler.inverse_transform.called
        assert result.y_pred is not None


class TestTrain:
    """Tests for train method."""

    @pytest.fixture
    def sample_csv_data(self):
        """Create sample CSV data in BytesIO."""
        csv_content = """var1,var2,var3,target
1.0,2.0,3.0,10.0
2.0,3.0,4.0,15.0
3.0,4.0,5.0,20.0
4.0,5.0,6.0,25.0
5.0,6.0,7.0,30.0
6.0,7.0,8.0,35.0
7.0,8.0,9.0,40.0
8.0,9.0,10.0,45.0
9.0,10.0,11.0,50.0
10.0,11.0,12.0,55.0
"""
        return BytesIO(csv_content.encode('utf-8'))

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_basic_workflow(
        self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
    ):
        """Test basic training workflow."""
        # Mock load_data to return a DataFrame
        mock_df = pd.DataFrame(
            {
                'var1': [1, 2, 3, 4, 5],
                'var2': [2, 3, 4, 5, 6],
                'var3': [3, 4, 5, 6, 7],
                'target': [10, 15, 20, 25, 30],
            }
        )
        mock_load_data.return_value = mock_df

        # Mock split_train_test to return train/test splits
        x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
        x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
        y_train = pd.Series([10, 15, 20], name='target')
        y_test = pd.Series([25, 30], name='target')
        mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)

        result = training_repo.train(sample_csv_data, sample_params)

        # Verify result structure
        assert result is not None
        assert isinstance(result, TrainModelResult)
        assert result.regr is not None
        assert result.x_train is not None
        assert result.x_test is not None
        assert result.y_train is not None
        assert result.y_test is not None
        assert result.process_data is not None
        assert result.scaler_dict is not None

        # Verify load_data was called correctly
        mock_load_data.assert_called_once_with(
            sample_csv_data, sample_params.line_separator, sample_params.decimal_separator
        )

        # Verify split was called
        assert mock_split_train_test.called

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_with_scaler(
        self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
    ):
        """Test training with scaler enabled."""
        sample_params.use_scaler = True

        mock_df = pd.DataFrame(
            {
                'var1': [1, 2, 3, 4, 5],
                'var2': [2, 3, 4, 5, 6],
                'var3': [3, 4, 5, 6, 7],
                'target': [10, 15, 20, 25, 30],
            }
        )
        mock_load_data.return_value = mock_df

        x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
        x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
        y_train = pd.Series([10, 15, 20], name='target')
        y_test = pd.Series([25, 30], name='target')
        mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)

        result = training_repo.train(sample_csv_data, sample_params)

        assert result is not None
        assert result.scaler_dict is not None

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_with_shuffle_enabled(
        self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
    ):
        """Test training with shuffle enabled."""
        sample_params.shuffle = True

        mock_df = pd.DataFrame(
            {
                'var1': [1, 2, 3, 4, 5],
                'var2': [2, 3, 4, 5, 6],
                'var3': [3, 4, 5, 6, 7],
                'target': [10, 15, 20, 25, 30],
            }
        )
        mock_load_data.return_value = mock_df

        x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
        x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
        y_train = pd.Series([10, 15, 20], name='target')
        y_test = pd.Series([25, 30], name='target')
        mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)

        training_repo.train(sample_csv_data, sample_params)

        # Verify split was called with shuffle=True
        call_kwargs = mock_split_train_test.call_args[1]
        assert call_kwargs['shuffle'] is True

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_with_different_train_size(
        self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
    ):
        """Test training with different train size."""
        sample_params.train_size = 70

        mock_df = pd.DataFrame(
            {
                'var1': [1, 2, 3, 4, 5],
                'var2': [2, 3, 4, 5, 6],
                'var3': [3, 4, 5, 6, 7],
                'target': [10, 15, 20, 25, 30],
            }
        )
        mock_load_data.return_value = mock_df

        x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
        x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
        y_train = pd.Series([10, 15, 20], name='target')
        y_test = pd.Series([25, 30], name='target')
        mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)

        training_repo.train(sample_csv_data, sample_params)

        # Verify split was called with train_size=0.7
        call_kwargs = mock_split_train_test.call_args[1]
        assert call_kwargs['train_size'] == 0.7

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_raises_on_empty_data_after_transform(
        self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
    ):
        """Test that ValueError is raised when transformed data is empty."""
        # Mock load_data to return empty DataFrame
        mock_df = pd.DataFrame(
            {
                'var1': [],
                'var2': [],
                'var3': [],
                'target': [],
            }
        )
        mock_load_data.return_value = mock_df

        with pytest.raises(ValueError, match='Data view is empty after transformation'):
            training_repo.train(sample_csv_data, sample_params)

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_logs_success(
        self,
        mock_load_data,
        mock_split_train_test,
        training_repo,
        mock_logger,
        sample_params,
        sample_csv_data,
    ):
        """Test that training success is logged."""
        mock_df = pd.DataFrame(
            {
                'var1': [1, 2, 3, 4, 5],
                'var2': [2, 3, 4, 5, 6],
                'var3': [3, 4, 5, 6, 7],
                'target': [10, 15, 20, 25, 30],
            }
        )
        mock_load_data.return_value = mock_df

        x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
        x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
        y_train = pd.Series([10, 15, 20], name='target')
        y_test = pd.Series([25, 30], name='target')
        mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)

        training_repo.train(sample_csv_data, sample_params)

        # Verify success was logged
        mock_logger.info.assert_called()
        assert any(
            'Model trained successfully' in str(call) for call in mock_logger.info.call_args_list
        )

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_with_custom_separators(
        self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
    ):
        """Test training with custom line and decimal separators."""
        sample_params.line_separator = ';'
        sample_params.decimal_separator = ','

        mock_df = pd.DataFrame(
            {
                'var1': [1, 2, 3, 4, 5],
                'var2': [2, 3, 4, 5, 6],
                'var3': [3, 4, 5, 6, 7],
                'target': [10, 15, 20, 25, 30],
            }
        )
        mock_load_data.return_value = mock_df

        x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
        x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
        y_train = pd.Series([10, 15, 20], name='target')
        y_test = pd.Series([25, 30], name='target')
        mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)

        training_repo.train(sample_csv_data, sample_params)

        # Verify load_data was called with custom separators
        mock_load_data.assert_called_once_with(sample_csv_data, ';', ',')

    @patch('model_manager.utils.repository.training_repository.split_train_test')
    @patch('model_manager.utils.repository.training_repository.load_data')
    def test_train_result_contains_all_fields(
        self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data
    ):
        """Test that TrainModelResult contains all expected fields."""
        mock_df = pd.DataFrame(
            {
                'var1': [1, 2, 3, 4, 5],
                'var2': [2, 3, 4, 5, 6],
                'var3': [3, 4, 5, 6, 7],
                'target': [10, 15, 20, 25, 30],
            }
        )
        mock_load_data.return_value = mock_df

        x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]})
        x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]})
        y_train = pd.Series([10, 15, 20], name='target')
        y_test = pd.Series([25, 30], name='target')
        mock_split_train_test.return_value = (x_train, x_test, y_train, y_test)

        result = training_repo.train(sample_csv_data, sample_params)

        # Verify all expected fields are present
        assert hasattr(result, 'params')
        assert hasattr(result, 'process_data')
        assert hasattr(result, 'x_train')
        assert hasattr(result, 'x_test')
        assert hasattr(result, 'y_train')
        assert hasattr(result, 'y_test')
        assert hasattr(result, 'regr')
        assert hasattr(result, 'scaler_dict')
        assert result.params == sample_params