"""Unit tests for sientia models module."""

from unittest.mock import MagicMock, patch

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

from model_manager.sientia.models import (
    ALLOWED_FRONTEND_DATE_FORMATS,
    FRONTEND_DATE_FORMAT_TO_STRFTIME,
    DataPreprocessor,
    LinearRegressionModel,
    _frontend_date_format_to_strftime,
    validate_frontend_date_format,
)


class TestFrontendDateFormatToStrftime:
    """Tests for _frontend_date_format_to_strftime (models module)."""

    def test_none_or_empty_returns_none(self):
        """None or empty string returns None."""
        assert _frontend_date_format_to_strftime(None) is None
        assert _frontend_date_format_to_strftime('') is None

    def test_all_six_allowed_formats_convert_correctly(self):
        """All allowed frontend formats map to expected strftime."""
        for frontend_fmt, strftime_fmt in FRONTEND_DATE_FORMAT_TO_STRFTIME.items():
            assert _frontend_date_format_to_strftime(frontend_fmt) == strftime_fmt

    def test_invalid_format_raises(self):
        """Invalid format raises ValueError with allowed list in message."""
        with raises(ValueError, match='Invalid date_format'):
            _frontend_date_format_to_strftime('yyyy-MM-dd')
        with raises(ValueError, match='Allowed formats'):
            _frontend_date_format_to_strftime('invalid')


class TestValidateFrontendDateFormat:
    """Tests for validate_frontend_date_format."""

    def test_none_or_empty_does_not_raise(self):
        """None or empty string does not raise."""
        validate_frontend_date_format(None)
        validate_frontend_date_format('')
        validate_frontend_date_format('   ')

    def test_allowed_formats_do_not_raise(self):
        """All allowed formats pass validation."""
        for fmt in ALLOWED_FRONTEND_DATE_FORMATS:
            validate_frontend_date_format(fmt)

    def test_invalid_format_raises(self):
        """Invalid format raises ValueError."""
        with raises(ValueError, match='Invalid date_format'):
            validate_frontend_date_format('yyyy-MM-dd')
        with raises(ValueError, match='Invalid date_format'):
            validate_frontend_date_format('custom-bad-format')


class _IterableWithContains:
    def __init__(self, iterable, contains_values):
        self._iterable = iterable
        self._contains = set(contains_values)

    def __iter__(self):
        return iter(self._iterable)

    def __contains__(self, item):
        return item in self._contains


# LinearRegressionModel Tests


def test_linear_regression_model_init_default():
    """Test LinearRegressionModel initialization with default parameters."""
    model = LinearRegressionModel()

    assert model.target_variable == ''
    assert model.variable_columns is None
    assert model.model_params is None
    assert model.clipping is None
    assert model.weights is None
    assert model.q1_target is None
    assert model.q3_target is None


def test_linear_regression_model_init_with_params():
    """Test LinearRegressionModel initialization with parameters."""
    target = 'target'
    variables = ['var1', 'var2']
    params = {'fit_intercept': True}
    clipping = {'min': 0, 'max': 100}
    weights = {'var1': 0.5, 'var2': 0.3}

    model = LinearRegressionModel(
        target_variable=target,
        variable_columns=variables,
        model_params=params,
        clipping=clipping,
        weights=weights,
    )

    assert model.target_variable == target
    assert model.variable_columns == variables
    assert model.model_params == params
    assert model.clipping == clipping
    assert model.weights == weights


def test_linear_regression_model_fit():
    """Test LinearRegressionModel fit method."""
    model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2'])

    data = pd.DataFrame(
        {'var1': [1, 2, 3, 4, 5], 'var2': [2, 3, 4, 5, 6], 'target': [3, 5, 7, 9, 11]}
    )

    result = model.fit(data)

    assert result is model
    assert model.q1_target is not None
    assert model.q3_target is not None
    assert model.weights is not None
    assert 'Bias' in model.weights


def test_linear_regression_model_fit_without_variable_columns():
    """Test LinearRegressionModel fit infers variable_columns when not set."""
    model = LinearRegressionModel(target_variable='target')

    data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})

    # Model should infer variable_columns from data (all columns except target)
    result = model.fit(data)
    assert result is model
    assert model.variable_columns == ['var1']


def test_linear_regression_model_predict_without_clipping():
    """Test LinearRegressionModel predict without clipping."""
    model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2'])

    train_data = pd.DataFrame(
        {'var1': [1, 2, 3, 4, 5], 'var2': [2, 3, 4, 5, 6], 'target': [3, 5, 7, 9, 11]}
    )
    model.fit(train_data)

    test_data = pd.DataFrame({'var1': [6, 7], 'var2': [7, 8]})
    predictions = model.predict(test_data)

    assert isinstance(predictions, np.ndarray)
    assert len(predictions) == 2


def test_linear_regression_model_predict_with_clipping_max():
    """Test LinearRegressionModel predict with clipping max."""
    model = LinearRegressionModel(
        target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 5}
    )

    train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]})
    model.fit(train_data)

    test_data = pd.DataFrame({'var1': [10]})
    predictions = model.predict(test_data)

    assert predictions[0] == model.q3_target


def test_linear_regression_model_predict_with_clipping_min():
    """Test LinearRegressionModel predict with clipping min."""
    model = LinearRegressionModel(
        target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 10}
    )

    train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]})
    model.fit(train_data)

    test_data = pd.DataFrame({'var1': [-10]})
    predictions = model.predict(test_data)

    assert predictions[0] == model.q1_target


def test_linear_regression_model_predict_with_clipping_within_range():
    """Test LinearRegressionModel predict with clipping but value within range."""
    model = LinearRegressionModel(
        target_variable='target', variable_columns=['var1'], clipping={'min': 0, 'max': 10}
    )

    train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 2, 3, 4, 5]})
    model.fit(train_data)

    test_data = pd.DataFrame({'var1': [3]})
    predictions = model.predict(test_data)

    # Prediction should be within range and not clipped
    assert 0 <= predictions[0] <= 10


# DataPreprocessor Tests


def test_data_preprocessor_init_default():
    """Test DataPreprocessor initialization with default parameters."""
    preprocessor = DataPreprocessor()

    assert preprocessor.date_column == ''
    assert preprocessor.target_variable == ''
    assert preprocessor.input_columns is None
    assert preprocessor.nan_treatment is None
    assert preprocessor.lag_train == {}
    assert preprocessor.lag_transform == {}
    assert preprocessor.scaler is None


def test_data_preprocessor_init_with_standard_scaler():
    """Test DataPreprocessor initialization with Standard Scaler."""
    preprocessor = DataPreprocessor(scaler_name='Standard Scaler')

    assert preprocessor.scaler is not None


def test_data_preprocessor_init_with_none_scaler():
    """Test DataPreprocessor initialization with None scaler."""
    preprocessor = DataPreprocessor(scaler_name='None')

    assert preprocessor.scaler is None


def test_data_preprocessor_init_with_unknown_scaler():
    """Test DataPreprocessor initialization with unknown scaler."""
    preprocessor = DataPreprocessor(scaler_name='Unknown')

    assert preprocessor.scaler is None


def test_data_preprocessor_init_with_custom_steps_order():
    """Test DataPreprocessor initialization with custom steps order."""
    custom_steps = ['Normalization', 'Feature Creation']
    preprocessor = DataPreprocessor(steps_order=custom_steps)

    assert 'Normalization' in preprocessor.steps_order
    assert 'Feature Creation' in preprocessor.steps_order
    assert len(preprocessor.steps_order) == 8  # Now includes RANGE_SELECTION step


def test_data_preprocessor_get_scaler():
    """Test DataPreprocessor get_scaler method."""
    preprocessor = DataPreprocessor(scaler_name='Standard Scaler')

    scaler = preprocessor.get_scaler()

    assert scaler is not None


@patch('model_manager.sientia.models.treat_nan')
def test_data_preprocessor_treat_discontinuities_with_treatment(mock_treat_nan):
    """Test treat_discontinuities with nan_treatment."""
    preprocessor = DataPreprocessor(nan_treatment='drop')
    data = pd.DataFrame({'col1': [1, 2, np.nan]})
    expected_data = pd.DataFrame({'col1': [1, 2]})
    mock_treat_nan.return_value = expected_data

    result = preprocessor.treat_discontinuities(data)

    mock_treat_nan.assert_called_once_with(data, 'drop')
    pd.testing.assert_frame_equal(result, expected_data)


def test_data_preprocessor_treat_discontinuities_without_treatment():
    """Test treat_discontinuities without nan_treatment."""
    preprocessor = DataPreprocessor()
    data = pd.DataFrame({'col1': [1, 2, 3]})

    result = preprocessor.treat_discontinuities(data)

    pd.testing.assert_frame_equal(result, data)


def test_data_preprocessor_lag_selection_with_lag():
    """Test lag_selection with lag."""
    preprocessor = DataPreprocessor()
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]})
    lag_dict = {'var1': 1}

    result = preprocessor.lag_selection(data, lag_dict)

    assert len(result) == 4
    assert result['var1'].iloc[0] == 1


def test_data_preprocessor_lag_selection_without_lag():
    """Test lag_selection without lag."""
    preprocessor = DataPreprocessor()
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]})
    lag_dict = {}

    result = preprocessor.lag_selection(data, lag_dict)

    assert len(result) == 5


def test_data_preprocessor_lag_selection_with_zero_lag():
    """Test lag_selection with zero lag."""
    preprocessor = DataPreprocessor()
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]})
    lag_dict = {'var1': 0}

    result = preprocessor.lag_selection(data, lag_dict)

    assert len(result) == 5


@patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer')
def test_data_preprocessor_treat_static_windows(mock_analyzer_class):
    """Test treat_static_windows method."""
    preprocessor = DataPreprocessor(static_threshold=3)
    data = pd.DataFrame({'col1': [1, 1, 1, 2, 3]})

    mock_analyzer = MagicMock()
    mock_analyzer_class.return_value = mock_analyzer
    mock_analyzer.get_treated_data.return_value = data

    preprocessor.treat_static_windows(data)

    mock_analyzer.infer_frequency.assert_called_once()
    assert mock_analyzer.identify_static_windows.called
    assert mock_analyzer.treat_static_windows.called


def test_data_preprocessor_treat_static_windows_without_threshold():
    """Test treat_static_windows without threshold."""
    preprocessor = DataPreprocessor()
    data = pd.DataFrame({'col1': [1, 2, 3]})

    result = preprocessor.treat_static_windows(data)

    pd.testing.assert_frame_equal(result, data)


@patch('model_manager.sientia.models.limit_dataset')
def test_data_preprocessor_adjust_limits(mock_limit_dataset):
    """Test adjust_limits method."""
    preprocessor = DataPreprocessor(low_lim={'col1': 0}, upp_lim={'col1': 10})
    data = pd.DataFrame({'col1': [1, 2, 3]})
    expected_data = pd.DataFrame({'col1': [1, 2, 3]})
    mock_limit_dataset.return_value = (expected_data, {'col1': 0}, {'col1': 10})

    result = preprocessor.adjust_limits(data)

    mock_limit_dataset.assert_called_once()
    pd.testing.assert_frame_equal(result, expected_data)


@patch('model_manager.sientia.models.create_features')
def test_data_preprocessor_create_features(mock_create_features):
    """Test create_features method."""
    preprocessor = DataPreprocessor(
        self_operations=['{var1}_{pow}_{2}'], cross_operations=['{var1}_{*}_{var2}']
    )
    data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4]})
    expected_data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var1_pow_2': [1, 4, 9]})
    mock_create_features.return_value = expected_data

    result = preprocessor.create_features(data)

    mock_create_features.assert_called_once()
    pd.testing.assert_frame_equal(result, expected_data)


def test_data_preprocessor_create_ar():
    """Test create_ar method."""
    preprocessor = DataPreprocessor(target_variable='target', ar_var='ar_target')
    data = pd.DataFrame({'target': [1, 2, 3, 4, 5]})

    result = preprocessor.create_ar(data)

    assert 'ar_target' in result.columns
    assert len(result) == 4


def test_data_preprocessor_create_ar_without_ar_var():
    """Test create_ar without ar_var."""
    preprocessor = DataPreprocessor(target_variable='target')
    data = pd.DataFrame({'target': [1, 2, 3, 4, 5]})

    result = preprocessor.create_ar(data)

    assert len(result) == 5


def test_data_preprocessor_create_lags():
    """Test create_lags method."""
    preprocessor = DataPreprocessor(created_lags={'var1': 1})
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]})

    result = preprocessor.create_lags(data)

    assert 'var1_lag1' in result.columns
    assert len(result) == 4


def test_data_preprocessor_create_lags_with_zero_lag():
    """Test create_lags with zero lag."""
    preprocessor = DataPreprocessor(created_lags={'var1': 0})
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]})

    result = preprocessor.create_lags(data)

    assert 'var1_lag0' not in result.columns
    assert len(result) == 5


def test_data_preprocessor_create_lags_without_created_lags():
    """Test create_lags without created_lags."""
    preprocessor = DataPreprocessor()
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]})

    result = preprocessor.create_lags(data)

    assert len(result) == 5


def test_data_preprocessor_create_lags_with_missing_column():
    """Test create_lags with missing column."""
    preprocessor = DataPreprocessor(created_lags={'var2': 1})
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5]})

    result = preprocessor.create_lags(data)

    assert 'var2_lag1' not in result.columns


def test_data_preprocessor_fit_with_x_and_y():
    """Test fit method with x and y."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4]})
    y = pd.Series([3, 5, 7], name='target')

    result = preprocessor.fit(x, y)

    assert result is preprocessor


def test_data_preprocessor_fit_with_only_x():
    """Test fit method with only x."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})

    result = preprocessor.fit(x)

    assert result is preprocessor


def test_data_preprocessor_fit_without_data():
    """Test fit method without data."""
    preprocessor = DataPreprocessor(target_variable='target', input_columns=['var1'])

    with raises(ValueError, match='No data was provided'):
        preprocessor.fit(None, None)


def test_data_preprocessor_fit_without_input_columns():
    """Test fit method without input_columns."""
    preprocessor = DataPreprocessor(target_variable='target')
    x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})

    with raises(AssertionError, match='input_columns must be set'):
        preprocessor.fit(x)


def test_data_preprocessor_fit_with_normalization():
    """Test fit method with normalization."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        scaler_name='Standard Scaler',
        scaler_params={},
        steps_order=['Normalization'],
    )
    x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})

    result = preprocessor.fit(x)

    assert result is preprocessor
    assert preprocessor.scaler_params is not None


def test_data_preprocessor_transform_with_timestamp():
    """Test transform method with timestamp column."""
    preprocessor = DataPreprocessor(
        target_variable='target', input_columns=['var1'], steps_order=['Discontinuity Treatment']
    )
    x = pd.DataFrame({'timestamp': [1, 2, 3], 'var1': [1, 2, 3], 'target': [3, 5, 7]})

    result = preprocessor.transform(x)

    assert 'timestamp' not in result.columns


def test_data_preprocessor_transform_without_timestamp():
    """Test transform method without timestamp column."""
    preprocessor = DataPreprocessor(
        target_variable='target', input_columns=['var1'], steps_order=['Discontinuity Treatment']
    )
    x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})

    result = preprocessor.transform(x)

    assert 'var1' in result.columns


def test_data_preprocessor_transform_without_input_columns():
    """Test transform method without input_columns."""
    preprocessor = DataPreprocessor(target_variable='target')
    x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})

    with raises(AssertionError, match='input_columns must be set'):
        preprocessor.transform(x)


def test_data_preprocessor_transform_with_feature_creation():
    """Test transform method with feature creation."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1'],
        self_operations=['{var1}_{pow}_{2}'],
        steps_order=['Feature Creation'],
    )
    x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})

    with patch('model_manager.sientia.models.create_features') as mock_create:
        mock_create.return_value = x
        preprocessor.transform(x)
        mock_create.assert_called_once()


def test_data_preprocessor_transform_with_lag_creation():
    """Test transform method with lag creation."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'ar_target'],
        ar_var='ar_target',
        created_lags={'var1': 1},
        steps_order=['Lag Creation'],
    )
    x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]})

    result = preprocessor.transform(x)

    assert 'ar_target' in result.columns
    assert 'var1_lag1' in result.columns


def test_data_preprocessor_transform_with_normalization():
    """Test transform method with normalization."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1'],
        scaler_name='Standard Scaler',
        scaler_params={},
        steps_order=['Normalization'],
    )
    train_x = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})
    preprocessor.fit(train_x)

    test_x = pd.DataFrame({'var1': [4, 5, 6], 'target': [9, 11, 13]})
    result = preprocessor.transform(test_x)

    assert 'var1' in result.columns


def test_data_preprocessor_get_required_columns_with_self_operations():
    """Test get_required_columns with self_operations."""
    preprocessor = DataPreprocessor(self_operations=['{var1}_{pow}_{2}'])
    existing_columns = ['var1', 'var2']

    result = preprocessor.get_required_columns(existing_columns)

    assert 'var1' in result


def test_data_preprocessor_get_required_columns_with_cross_operations():
    """Test get_required_columns with cross_operations."""
    preprocessor = DataPreprocessor(cross_operations=['{var1}_{*}_{var2}'])
    existing_columns = ['var1', 'var2']

    result = preprocessor.get_required_columns(existing_columns)

    assert 'var1' in result
    assert 'var2' in result


def test_data_preprocessor_get_required_columns_with_created_lags():
    """Test get_required_columns with created_lags."""
    preprocessor = DataPreprocessor(created_lags={'var1': 1})
    existing_columns = ['var1', 'var2']

    result = preprocessor.get_required_columns(existing_columns)

    assert 'var1' in result


def test_data_preprocessor_get_required_columns_with_missing_columns():
    """Test get_required_columns with missing columns in existing_columns."""
    preprocessor = DataPreprocessor(
        self_operations=['{var3}_{pow}_{2}'], cross_operations=['{var4}_{*}_{var5}']
    )
    existing_columns = ['var1', 'var2']

    result = preprocessor.get_required_columns(existing_columns)

    assert 'var3' in result
    assert 'var4' in result
    assert 'var5' in result


def test_data_preprocessor_get_required_columns_removes_duplicates():
    """Test get_required_columns removes duplicates from self_operations."""
    preprocessor = DataPreprocessor(self_operations=['{var1}_{pow}_{2}'], created_lags={'var1': 1})
    existing_columns = ['var1', 'var2']

    result = preprocessor.get_required_columns(existing_columns)

    # var1 is in existing_columns, so it should not be in required_columns
    assert 'var1' not in result or result.count('var1') <= 1


def test_data_preprocessor_get_required_columns_removes_self_operations_branch():
    """Ensure line 278 removes columns present in self_operations iterable."""
    preprocessor = DataPreprocessor(
        self_operations=_IterableWithContains(['{var1}_{pow}_{2}'], contains_values=['var1'])
    )
    existing_columns: list[str] = []

    result = preprocessor.get_required_columns(existing_columns)

    assert 'var1' not in result


def test_data_preprocessor_get_required_columns_removes_cross_operations_branch():
    """Ensure line 285 removes columns present in cross_operations iterable."""
    preprocessor = DataPreprocessor(
        cross_operations=_IterableWithContains(['{var1}_{*}_{var2}'], contains_values=['var1'])
    )
    existing_columns: list[str] = []

    result = preprocessor.get_required_columns(existing_columns)

    assert 'var1' not in result
    assert 'var2' in result


def test_data_preprocessor_get_required_columns_removes_created_lags():
    """Test get_required_columns removes columns from created_lags when column is in created_lags dict - covers line 292."""
    preprocessor = DataPreprocessor(created_lags={'var1': 1, 'var2': 1})
    existing_columns = ['var3']

    result = preprocessor.get_required_columns(existing_columns)

    # var1 and var2 should be removed because they're in created_lags dict and not in existing_columns
    assert 'var1' not in result
    assert 'var2' not in result


def test_data_preprocessor_fit_all_steps():
    """Test fit method with all steps."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1'],
        nan_treatment='drop',
        lag_train={'var1': 1},
        static_threshold=3,
        low_lim={'var1': 0},
        upp_lim={'var1': 10},
        scaler_name='Standard Scaler',
        scaler_params={},
    )
    x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]})

    with patch('model_manager.sientia.models.treat_nan') as mock_treat:
        with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'):
            with patch('model_manager.sientia.models.limit_dataset') as mock_limit:
                mock_treat.return_value = x
                mock_limit.return_value = (x, {'var1': 0}, {'var1': 10})
                result = preprocessor.fit(x)

    assert result is preprocessor


def test_data_preprocessor_transform_all_steps():
    """Test transform method with all steps."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1'],
        nan_treatment='drop',
        lag_transform={'var1': 1},
        static_threshold=3,
        low_lim={'var1': 0},
        upp_lim={'var1': 10},
        scaler_name='Standard Scaler',
        scaler_params={},
        self_operations=['{var1}_{pow}_{2}'],
        ar_var='ar_target',
        created_lags={'var1': 1},
    )
    train_x = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]})

    with patch('model_manager.sientia.models.treat_nan') as mock_treat:
        with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'):
            with patch('model_manager.sientia.models.limit_dataset') as mock_limit:
                with patch('model_manager.sientia.models.create_features') as mock_create:
                    mock_treat.return_value = train_x
                    mock_limit.return_value = (train_x, {'var1': 0}, {'var1': 10})
                    mock_create.return_value = train_x
                    preprocessor.fit(train_x)

    test_x = pd.DataFrame({'var1': [6, 7, 8, 9, 10], 'target': [13, 15, 17, 19, 21]})

    with patch('model_manager.sientia.models.treat_nan') as mock_treat:
        with patch('model_manager.sientia.models.TimeSeriesDiscontinuityAnalyzer'):
            with patch('model_manager.sientia.models.limit_dataset') as mock_limit:
                with patch('model_manager.sientia.models.create_features') as mock_create:
                    mock_treat.return_value = test_x
                    mock_limit.return_value = (test_x, {'var1': 0}, {'var1': 10})
                    mock_create.return_value = test_x
                    result = preprocessor.transform(test_x)

    assert isinstance(result, pd.DataFrame)


# ============================================================================
# Additional tests for coverage - LinearRegressionModel
# ============================================================================


def test_linear_regression_model_fit_without_target_variable():
    """Test fit raises error when target_variable is not set."""
    model = LinearRegressionModel()
    data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})

    with raises(ValueError, match='target_variable must be set before fitting'):
        model.fit(data)


def test_linear_regression_model_fit_with_missing_columns():
    """Test fit raises error when variable_columns are missing from data."""
    model = LinearRegressionModel(
        target_variable='target', variable_columns=['var1', 'var_missing']
    )
    data = pd.DataFrame({'var1': [1, 2, 3], 'target': [3, 5, 7]})

    with raises(ValueError, match='Columns not found in input data'):
        model.fit(data)


def test_linear_regression_model_fit_with_missing_target():
    """Test fit raises error when target_variable is not in data."""
    model = LinearRegressionModel(target_variable='missing_target', variable_columns=['var1'])
    data = pd.DataFrame({'var1': [1, 2, 3], 'other': [3, 5, 7]})

    with raises(ValueError, match='Target variable missing_target not found in input data'):
        model.fit(data)


def test_linear_regression_model_fit_with_inf_values():
    """Test fit handles infinite values by converting to NaN."""
    model = LinearRegressionModel(target_variable='target', variable_columns=['var1', 'var2'])
    # Include some inf values that will be converted to NaN and rows dropped
    data = pd.DataFrame(
        {
            'var1': [1.0, 2.0, 3.0, np.inf, 5.0],
            'var2': [2.0, 3.0, 4.0, 5.0, 6.0],
            'target': [3.0, 5.0, 7.0, 9.0, 11.0],
        }
    )

    model.fit(data)

    # Model should fit successfully after removing row with inf
    assert model.weights is not None
    assert 'var1' in model.variable_columns
    assert 'var2' in model.variable_columns


def test_linear_regression_model_fit_polynomial():
    """Test fit with polynomial features."""
    model = LinearRegressionModel(target_variable='target', variable_columns=['var1'], degree=2)
    data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [3, 5, 7, 9, 11]})

    model.fit(data)

    assert model.poly_feature_names is not None
    assert len(model.poly_feature_names) > 1


def test_linear_regression_model_predict_polynomial():
    """Test predict with polynomial features."""
    model = LinearRegressionModel(target_variable='target', variable_columns=['var1'], degree=2)
    train_data = pd.DataFrame({'var1': [1, 2, 3, 4, 5], 'target': [1, 4, 9, 16, 25]})
    model.fit(train_data)

    test_data = pd.DataFrame({'var1': [6, 7]})
    predictions = model.predict(test_data)

    assert isinstance(predictions, np.ndarray)
    assert len(predictions) == 2


def test_linear_regression_model_create_poly_features_degree_1():
    """Test create_poly_features returns input unchanged when degree <= 1."""
    model = LinearRegressionModel(degree=1)
    data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]})

    result = model.create_poly_features(data, fit=True)

    pd.testing.assert_frame_equal(result, data)


def test_linear_regression_model_create_poly_features_not_fitted():
    """Test create_poly_features raises error when not fitted and fit=False."""
    model = LinearRegressionModel(degree=2)
    data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]})

    with raises(ValueError, match='PolynomialFeatures not fitted'):
        model.create_poly_features(data, fit=False)


def test_linear_regression_model_create_poly_features_transform():
    """Test create_poly_features with fit=False after fitting."""
    model = LinearRegressionModel(degree=2)
    train_data = pd.DataFrame({'var1': [1, 2, 3], 'var2': [4, 5, 6]})
    model.create_poly_features(train_data, fit=True)

    test_data = pd.DataFrame({'var1': [4, 5], 'var2': [7, 8]})
    result = model.create_poly_features(test_data, fit=False)

    assert isinstance(result, pd.DataFrame)
    assert len(result.columns) > 2


def test_linear_regression_model_get_regressor():
    """Test get_regressor returns the underlying LinearRegression model."""
    model = LinearRegressionModel()

    regressor = model.get_regressor()

    from sklearn.linear_model import LinearRegression

    assert isinstance(regressor, LinearRegression)


# ============================================================================
# Additional tests for coverage - DataPreprocessor
# ============================================================================


def test_data_preprocessor_parse_datetime_with_frontend_format():
    """When date_format is set, _parse_datetime uses strftime mapping (covers format branch)."""
    preprocessor = DataPreprocessor(date_format='dd/MM/yyyy HH:mm:ss')
    ts = preprocessor._parse_datetime('15/01/2024 10:30:00')
    assert ts is not None


@patch('model_manager.sientia.models.treat_nan')
def test_data_preprocessor_treat_discontinuities_linear_interpolation(mock_treat_nan):
    """Test treat_discontinuities with 'linear interpolation' treatment."""
    preprocessor = DataPreprocessor(nan_treatment='linear interpolation')
    data = pd.DataFrame({'col1': [1, np.nan, 3]})
    expected_data = pd.DataFrame({'col1': [1.0, 2.0, 3.0]})
    mock_treat_nan.return_value = expected_data

    result = preprocessor.treat_discontinuities(data)

    # Should map 'linear interpolation' to 'fill linear'
    mock_treat_nan.assert_called_once_with(data, 'fill linear')
    pd.testing.assert_frame_equal(result, expected_data)


def test_data_preprocessor_range_selection_with_start_date():
    """Test range_selection filters by start_date."""
    preprocessor = DataPreprocessor(start_date='2023-01-02')
    data = pd.DataFrame(
        {'col1': [1, 2, 3]},
        index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
    )

    result = preprocessor.range_selection(data)

    assert len(result) == 2
    assert result.index[0] == pd.Timestamp('2023-01-02')


def test_data_preprocessor_range_selection_with_end_date():
    """Test range_selection filters by end_date."""
    preprocessor = DataPreprocessor(end_date='2023-01-02')
    data = pd.DataFrame(
        {'col1': [1, 2, 3]},
        index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
    )

    result = preprocessor.range_selection(data)

    assert len(result) == 2
    assert result.index[-1] == pd.Timestamp('2023-01-02')


def test_data_preprocessor_range_selection_with_invalid_start_date():
    """Test range_selection handles invalid start_date gracefully."""
    preprocessor = DataPreprocessor(start_date='invalid-date')
    data = pd.DataFrame(
        {'col1': [1, 2, 3]},
        index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
    )

    result = preprocessor.range_selection(data)

    # Should skip filtering and return original data
    assert len(result) == 3


def test_data_preprocessor_range_selection_with_invalid_end_date():
    """Test range_selection handles invalid end_date gracefully."""
    preprocessor = DataPreprocessor(end_date='invalid-date')
    data = pd.DataFrame(
        {'col1': [1, 2, 3]},
        index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
    )

    result = preprocessor.range_selection(data)

    # Should skip filtering and return original data
    assert len(result) == 3


def test_data_preprocessor_range_selection_with_removed_intervals():
    """Test range_selection removes specified intervals."""
    preprocessor = DataPreprocessor(removed_intervals=[['2023-01-02', '2023-01-03']])
    data = pd.DataFrame(
        {'col1': [1, 2, 3, 4, 5]},
        index=pd.to_datetime(
            ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04', '2023-01-05']
        ),
    )

    result = preprocessor.range_selection(data)

    assert len(result) == 3
    assert pd.Timestamp('2023-01-02') not in result.index
    assert pd.Timestamp('2023-01-03') not in result.index


def test_data_preprocessor_range_selection_with_invalid_interval():
    """Test range_selection handles invalid interval dates gracefully."""
    preprocessor = DataPreprocessor(removed_intervals=[['invalid', 'dates']])
    data = pd.DataFrame(
        {'col1': [1, 2, 3]},
        index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
    )

    result = preprocessor.range_selection(data)

    # Should skip invalid interval and return original data
    assert len(result) == 3


def test_data_preprocessor_range_selection_with_short_interval():
    """Test range_selection skips intervals with less than 2 elements."""
    preprocessor = DataPreprocessor(removed_intervals=[['2023-01-02']])
    data = pd.DataFrame(
        {'col1': [1, 2, 3]},
        index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']),
    )

    result = preprocessor.range_selection(data)

    # Should skip short interval and return original data
    assert len(result) == 3


def test_data_preprocessor_predict():
    """Test predict method removes target and preserves feature order."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
    preprocessor.fit(train_x)

    test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'target': [9, 11]})
    result = preprocessor.predict(test_x)

    assert 'target' not in result.columns
    assert 'var1' in result.columns
    assert 'var2' in result.columns


def test_data_preprocessor_predict_without_target():
    """Test predict when target is not in transformed data."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
    preprocessor.fit(train_x)

    # Include target in test data so transform works, predict will remove it
    test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'target': [9, 11]})
    result = preprocessor.predict(test_x)

    # Target should be removed by predict
    assert 'target' not in result.columns
    assert 'var1' in result.columns
    assert 'var2' in result.columns


def test_data_preprocessor_predict_preserves_feature_order():
    """Test predict preserves feature order from fit."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
    preprocessor.fit(train_x)

    # Test data has columns in different order
    test_x = pd.DataFrame({'var2': [5, 6], 'var1': [4, 5], 'target': [9, 11]})
    result = preprocessor.predict(test_x)

    # Should have columns in same order as during fit
    assert list(result.columns) == ['var1', 'var2']


# ============================================================================
# Additional tests for 100% coverage
# ============================================================================


def test_data_preprocessor_predict_target_not_in_columns():
    """Test predict when target_variable is not in transformed data columns."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
    preprocessor.fit(train_x)

    # Test data without target column - predict should still work
    test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6]})
    # Need to add target for transform to work, then it gets removed
    test_x['target'] = [9, 11]

    # Manually remove target before calling predict to test the branch
    preprocessor_copy = DataPreprocessor(
        target_variable='nonexistent_target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    preprocessor_copy.fit(train_x.rename(columns={'target': 'nonexistent_target'}))

    test_x_no_target = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'nonexistent_target': [9, 11]})
    result = preprocessor_copy.predict(test_x_no_target)

    assert 'var1' in result.columns
    assert 'var2' in result.columns


def test_data_preprocessor_predict_without_fitted_feature_order():
    """Test predict when _fitted_feature_order is None."""
    preprocessor = DataPreprocessor(
        target_variable='target',
        input_columns=['var1', 'var2'],
        steps_order=['Discontinuity Treatment'],
    )
    train_x = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'target': [3, 5, 7]})
    preprocessor.fit(train_x)

    # Manually set _fitted_feature_order to None to test the branch
    preprocessor._fitted_feature_order = None

    test_x = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'target': [9, 11]})
    result = preprocessor.predict(test_x)

    # Should still work, just without reordering
    assert 'target' not in result.columns
    assert 'var1' in result.columns
    assert 'var2' in result.columns