"""Unit tests for TrainingRepository."""

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

import numpy as np
import pandas as pd
from pytest import fixture, raises

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


@fixture
def logger():
    """Create a mock logger."""
    return MagicMock()


@fixture
def training_repository(logger):
    """Create a TrainingRepository instance."""
    return TrainingRepository(logger)


@fixture
def train_params():
    """Create sample training parameters."""
    return TrainModelParams(
        variable_columns=['feature1', 'feature2'],
        lag_train=1,
        lag_val=1,
        target_variable='target',
        rem_static_win=False,
        low_lim={'feature1': 0.0, 'feature2': 0.0},
        upp_lim={'feature1': 100.0, 'feature2': 100.0},
        window=10,
        use_scaler=True,
        include_ar=False,
        bucket_name='test-bucket',
        file_name='test.csv',
        line_separator='\n',
        decimal_separator='.',
        train_size=80,
        shuffle=True,
        experiment_run_id=123,
        experiment_name='test_experiment',
        removed_intervals=[],
    )


@fixture
def sample_csv_data():
    """Create sample CSV data."""
    csv_content = """feature1,feature2,target
1.0,2.0,10.0
2.0,3.0,15.0
3.0,4.0,20.0
4.0,5.0,25.0
5.0,6.0,30.0
6.0,7.0,35.0
7.0,8.0,40.0
8.0,9.0,45.0
9.0,10.0,50.0
10.0,11.0,55.0
"""
    return BytesIO(csv_content.encode())


@patch('model_manager.utils.repository.training_repository.load_data')
@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
@patch('model_manager.utils.repository.training_repository.split_train_test')
@patch('model_manager.utils.repository.training_repository.LinearRegressionModel')
def test_train_success(
    mock_linear_model,
    mock_split,
    mock_preprocessor_class,
    mock_load_data,
    training_repository,
    train_params,
    sample_csv_data,
):
    """Test successful model training."""
    # Setup mocks
    mock_data = pd.DataFrame(
        {'feature1': [1, 2, 3, 4, 5], 'feature2': [2, 3, 4, 5, 6], 'target': [10, 15, 20, 25, 30]}
    )
    mock_load_data.return_value = mock_data

    mock_preprocessor = MagicMock()
    mock_preprocessor_class.return_value = mock_preprocessor
    mock_preprocessor.transform.return_value = mock_data

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

    mock_model = MagicMock()
    mock_linear_model.return_value = mock_model

    mock_scaler = MagicMock()
    mock_preprocessor.get_scaler.return_value = mock_scaler

    # Execute
    result = training_repository.train(sample_csv_data, train_params)

    # Assertions
    assert isinstance(result, TrainModelResult)
    assert result.params == train_params
    assert result.process_data == mock_preprocessor
    assert result.regr == mock_model
    mock_load_data.assert_called_once()
    mock_preprocessor.fit.assert_called_once()
    mock_model.fit.assert_called_once()


@patch('model_manager.utils.repository.training_repository.load_data')
def test_train_empty_data_after_transform(
    mock_load_data, training_repository, train_params, sample_csv_data
):
    """Test training with empty data after transformation."""
    mock_data = pd.DataFrame({'feature1': [], 'feature2': [], 'target': []})
    mock_load_data.return_value = mock_data

    with patch.object(training_repository, 'init_data_preprocessor') as mock_init:
        mock_preprocessor = MagicMock()
        mock_init.return_value = mock_preprocessor
        mock_preprocessor.transform.return_value = pd.DataFrame()

        with raises(ValueError, match='Data view is empty after transformation'):
            training_repository.train(sample_csv_data, train_params)


def test_init_scaler_dict_with_minmax_scaler(training_repository, train_params):
    """Test scaler dict initialization with MinMaxScaler."""
    mock_preprocessor = MagicMock()
    mock_scaler = MagicMock()
    mock_scaler.x_min = [0.0, 1.0]
    mock_scaler.x_max = [10.0, 11.0]
    mock_scaler.y_min = 5.0
    mock_scaler.y_max = 50.0
    mock_preprocessor.get_scaler.return_value = mock_scaler

    # Patch isinstance to return True for MinMaxScaler
    with patch(
        'model_manager.utils.repository.training_repository.isinstance',
        side_effect=lambda obj, cls: cls.__name__ == 'MinMaxScaler',
    ):
        result = training_repository.init_scaler_dict(mock_preprocessor, train_params)

    assert result is not None
    assert 'feature1' in result
    assert 'feature2' in result
    assert 'target' in result
    assert result['feature1'] == {'min': 0.0, 'max': 10.0}
    assert result['feature2'] == {'min': 1.0, 'max': 11.0}
    assert result['target'] == {'min': 5.0, 'max': 50.0}


def test_init_scaler_dict_with_z_scaler(training_repository, train_params):
    """Test scaler dict initialization with Z_Scaler."""
    mock_preprocessor = MagicMock()
    mock_scaler = MagicMock()
    mock_scaler.create_dict.return_value = {'mean': 5.0, 'std': 2.0}
    mock_preprocessor.get_scaler.return_value = mock_scaler

    # Patch isinstance to return True for Z_Scaler
    with patch(
        'model_manager.utils.repository.training_repository.isinstance',
        side_effect=lambda obj, cls: cls.__name__ == 'Z_Scaler',
    ):
        result = training_repository.init_scaler_dict(mock_preprocessor, train_params)

    assert result == {'mean': 5.0, 'std': 2.0}
    mock_scaler.create_dict.assert_called_once()


def test_init_scaler_dict_without_scaler(training_repository):
    """Test scaler dict initialization when use_scaler is False."""
    train_params_no_scaler = TrainModelParams(
        variable_columns=['feature1'],
        lag_train=1,
        lag_val=1,
        target_variable='target',
        rem_static_win=False,
        low_lim={'feature1': 0.0},
        upp_lim={'feature1': 100.0},
        window=10,
        use_scaler=False,
        include_ar=False,
        bucket_name='test',
        file_name='test.csv',
        line_separator='\n',
        decimal_separator='.',
        train_size=80,
        shuffle=True,
        experiment_run_id=123,
        experiment_name='test',
        removed_intervals=[],
    )

    mock_preprocessor = MagicMock()
    result = training_repository.init_scaler_dict(mock_preprocessor, train_params_no_scaler)

    assert result == {}


@patch('model_manager.utils.repository.training_repository.mse')
@patch('model_manager.utils.repository.training_repository.mae')
@patch('model_manager.utils.repository.training_repository.r2')
def test_after_train_calculation_with_scaler(
    mock_r2, mock_mae, mock_mse, training_repository, train_params
):
    """Test post-training calculations with scaler."""
    # Setup mock train result
    mock_train_result = MagicMock(spec=TrainModelResult)
    mock_train_result.params = train_params
    mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [2, 3, 4]})
    mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5], 'feature2': [5, 6]})
    mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
    mock_train_result.y_test = pd.Series([25, 30], name='target')

    mock_regr = MagicMock()
    mock_regr.predict.return_value = np.array([24.5, 29.5])
    mock_train_result.regr = mock_regr

    mock_scaler = MagicMock()
    mock_scaler.denormalize_single_input.side_effect = lambda x, col: x
    mock_scaler.denormalize_predictions.side_effect = lambda x, col: x

    mock_process_data = MagicMock()
    mock_process_data.get_scaler.return_value = mock_scaler
    mock_train_result.process_data = mock_process_data

    # Setup metric mocks
    mock_mse.return_value = 0.5
    mock_mae.return_value = 0.3
    mock_r2.return_value = 0.95

    # Execute
    result = training_repository.after_train_calculation(train_params, mock_train_result)

    # Assertions
    assert result == mock_train_result
    assert result.mse_val == 0.5
    assert result.mae_val == 0.3
    assert result.r2_val == 0.95
    assert result.y_pred is not None
    mock_regr.predict.assert_called_once()
    mock_mse.assert_called_once()
    mock_mae.assert_called_once()
    mock_r2.assert_called_once()


@patch('model_manager.utils.repository.training_repository.mse')
@patch('model_manager.utils.repository.training_repository.mae')
@patch('model_manager.utils.repository.training_repository.r2')
def test_after_train_calculation_without_scaler(mock_r2, mock_mae, mock_mse, training_repository):
    """Test post-training calculations without scaler."""
    train_params_no_scaler = TrainModelParams(
        variable_columns=['feature1'],
        lag_train=1,
        lag_val=1,
        target_variable='target',
        rem_static_win=False,
        low_lim={'feature1': 0.0},
        upp_lim={'feature1': 100.0},
        window=10,
        use_scaler=False,
        include_ar=False,
        bucket_name='test',
        file_name='test.csv',
        line_separator='\n',
        decimal_separator='.',
        train_size=80,
        shuffle=True,
        experiment_run_id=123,
        experiment_name='test',
        removed_intervals=[],
    )

    mock_train_result = MagicMock(spec=TrainModelResult)
    mock_train_result.params = train_params_no_scaler
    mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3]})
    mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5]})
    mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
    mock_train_result.y_test = pd.Series([25, 30], name='target')

    mock_regr = MagicMock()
    mock_regr.predict.return_value = np.array([24.5, 29.5])
    mock_train_result.regr = mock_regr

    # Setup metric mocks
    mock_mse.return_value = 0.5
    mock_mae.return_value = 0.3
    mock_r2.return_value = 0.95

    # Execute
    result = training_repository.after_train_calculation(train_params_no_scaler, mock_train_result)

    # Assertions
    assert result.mse_val == 0.5
    assert result.mae_val == 0.3
    assert result.r2_val == 0.95


@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
def test_init_data_preprocessor(mock_preprocessor_class, training_repository, train_params):
    """Test DataPreprocessor initialization."""
    mock_preprocessor = MagicMock()
    mock_preprocessor_class.return_value = mock_preprocessor

    result = training_repository.init_data_preprocessor(train_params)

    assert result == mock_preprocessor
    mock_preprocessor_class.assert_called_once()
    call_kwargs = mock_preprocessor_class.call_args[1]
    assert call_kwargs['target_variable'] == 'target'
    assert call_kwargs['input_columns'] == ['feature1', 'feature2']
    assert call_kwargs['low_lim'] == {'feature1': 0.0, 'feature2': 0.0}
    assert call_kwargs['upp_lim'] == {'feature1': 100.0, 'feature2': 100.0}