"""
Winnow Test Suite
Comprehensive unit tests
"""

import pytest
import numpy as np
import pandas as pd
import sys
import os

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from winnow import BaseSelector, AdvancedSelector
from winnow.base_selector import SelectionResult
from winnow.advanced_selector import AdvancedSelectionResult


# ============================================================================
# FIXTURES
# ============================================================================

@pytest.fixture
def sample_data():
    """Basic test data."""
    np.random.seed(42)
    n_samples = 100
    
    X = pd.DataFrame({
        'f1': np.random.randn(n_samples),
        'f2': np.random.randn(n_samples),
        'f3': np.random.randn(n_samples),
    })
    y = pd.Series(np.random.randint(0, 2, n_samples))
    
    return X, y


@pytest.fixture
def data_with_high_correlation():
    """Data containing highly correlated features."""
    np.random.seed(42)
    n_samples = 100
    
    f1 = np.random.randn(n_samples)
    f2 = f1 * 0.95 + np.random.randn(n_samples) * 0.05  # high correlation with f1
    f3 = np.random.randn(n_samples)
    
    X = pd.DataFrame({'f1': f1, 'f2': f2, 'f3': f3})
    y = pd.Series((f1 > 0).astype(int))
    
    return X, y


@pytest.fixture
def data_with_low_variance():
    """Data containing low variance features."""
    np.random.seed(42)
    n_samples = 100
    
    X = pd.DataFrame({
        'f1': np.random.randn(n_samples),
        'f2': np.ones(n_samples) * 5 + np.random.randn(n_samples) * 0.001,  # low variance
        'f3': np.random.randn(n_samples),
    })
    y = pd.Series(np.random.randint(0, 2, n_samples))
    
    return X, y


@pytest.fixture
def data_with_low_target_correlation():
    """Data containing features with low target correlation."""
    np.random.seed(42)
    n_samples = 100
    
    f1 = np.random.randn(n_samples)
    f2 = np.random.randn(n_samples)  # uncorrelated with target
    
    X = pd.DataFrame({'f1': f1, 'f2': f2})
    y = pd.Series((f1 > 0).astype(int))  # dependent only on f1
    
    return X, y


@pytest.fixture
def feature_importance():
    """Sample feature importance scores."""
    return {
        'f1': 0.5,
        'f2': 0.3,
        'f3': 0.2
    }


# ============================================================================
# BASE SELECTOR TESTS
# ============================================================================

class TestBaseSelector:
    """Unit tests for BaseSelector."""
    
    def test_initialization_default_params(self):
        """Test initialization with default parameters."""
        selector = BaseSelector()
        
        assert selector.correlation_threshold == 0.8
        assert selector.target_correlation_threshold == 0.05
        assert selector.variance_threshold == 0.01
        assert selector._is_fitted == False
    
    def test_fit(self, sample_data):
        """Test fit method."""
        X, y = sample_data
        selector = BaseSelector()
        
        result = selector.fit(X, y)
        
        assert result is selector  # Fluent interface
        assert selector._is_fitted == True
        assert selector._variances is not None
    
    def test_transform_without_fit_raises_error(self, sample_data):
        """Calling transform without fit should raise error."""
        X, y = sample_data
        selector = BaseSelector()
        
        with pytest.raises(ValueError, match="not fitted yet"):
            selector.transform(X)
    
    def test_removes_high_correlation_features(self, data_with_high_correlation):
        """Should remove highly correlated features."""
        X, y = data_with_high_correlation
        selector = BaseSelector(correlation_threshold=0.7)
        
        result = selector.fit_transform(X, y)
        
        assert len(result.removed_features['high_correlation']) > 0
        assert ('f1' in result.selected_features) ^ ('f2' in result.selected_features)
    
    def test_invalid_input_not_dataframe(self):
        """Should raise error if X is not a DataFrame."""
        selector = BaseSelector()
        with pytest.raises(TypeError, match="pandas DataFrame"):
            selector.fit([[1, 2]], pd.Series([0]))


# ============================================================================
# ADVANCED SELECTOR TESTS
# ============================================================================

class TestAdvancedSelector:
    """Unit tests for AdvancedSelector."""
    
    def test_pca_enabled(self, sample_data):
        """Should work when PCA is enabled."""
        X, y = sample_data
        selector = AdvancedSelector(use_pca=True, pca_variance_threshold=0.95)
        
        result = selector.fit_transform(X, y)
        
        assert result.pca_components is not None
        assert result.pca_explained_variance_ratio is not None

    def test_custom_selection_method(self, sample_data):
        """Should execute custom selection methods."""
        X, y = sample_data
        selector = AdvancedSelector()
        
        # Method that removes a specific feature
        def remove_f1(X, y, features):
            return ['f1'] if 'f1' in features else []
            
        selector.add_custom_method("remove_f1", remove_f1)
        result = selector.fit_transform(X, y)
        
        assert 'f1' not in result.selected_features
        assert 'f1' in result.removed_features['remove_f1']


# ============================================================================
# TARGET COLUMN TESTS
# ============================================================================

class TestTargetColumn:
    """Tests for target_column parameter."""
    
    @pytest.fixture
    def dataframe_with_target(self):
        """DataFrame containing target."""
        np.random.seed(42)
        n_samples = 100
        df = pd.DataFrame({
            'f1': np.random.randn(n_samples),
            'f2': np.random.randn(n_samples),
            'target': np.random.randint(0, 2, n_samples)
        })
        return df
    
    def test_target_column_parameter(self, dataframe_with_target):
        """Should work with target_column parameter."""
        df = dataframe_with_target
        selector = BaseSelector()
        
        result = selector.fit_transform(df, target_column='target')
        
        assert isinstance(result, SelectionResult)
        assert 'target' not in result.selected_features
    
    def test_target_not_found(self, dataframe_with_target):
        """Should raise error if target column is not found."""
        df = dataframe_with_target
        selector = BaseSelector()
        with pytest.raises(ValueError, match="not found"):
            selector.fit(df, target_column='nonexistent')


if __name__ == "__main__":
    pytest.main([__file__, "-v", "--tb=short"])
