# tests/test_enhanced_sampling.py
"""
Tests for the enhanced sampling functionality.
"""

import pytest
import pandas as pd
import numpy as np
from fmatch.core.sampling import (
    ReservoirSampler, 
    AdaptiveSampler, 
    EarlyAcceptChecker,
    stratified_domain_sample
)
from fmatch.core.utils import analyze_domain_matching_strategy, calculate_domain_weight


class TestReservoirSampler:
    def test_reservoir_sampler_maintains_size(self):
        """Test that reservoir maintains correct size."""
        sampler = ReservoirSampler(k=100)
        for i in range(1000):
            sampler.add(i)
        
        assert len(sampler.get_sample()) == 100
        
    def test_reservoir_sampler_distribution(self):
        """Test that reservoir sampling is unbiased."""
        sampler = ReservoirSampler(k=100, random_state=42)
        
        # Add numbers 0-999
        for i in range(1000):
            sampler.add(i)
        
        sample = sampler.get_sample()
        # Sample mean should be close to population mean (499.5)
        assert 400 < np.mean(sample) < 600
        
    def test_get_sample_series(self):
        """Test getting sample as pandas Series."""
        sampler = ReservoirSampler(k=10)
        for i in range(20):
            sampler.add(f"item_{i}")
        
        series = sampler.get_sample_series("test_col")
        assert isinstance(series, pd.Series)
        assert series.name == "test_col"
        assert len(series) == 10


class TestAdaptiveSampler:
    def test_adaptive_sampler_converges(self):
        """Test that adaptive sampler converges on stable metric."""
        # Create data with known 50% split
        data = pd.Series(['A'] * 500 + ['B'] * 500)
        sampler = AdaptiveSampler(target_se=0.02)
        
        def split_metric(samples):
            unique_vals = set(samples)
            return len([s for s in samples if s == 'A']) / len(samples)
        
        samples, metric, se, info = sampler.sample_until_stable(data, split_metric)
        
        assert info['converged'] is True
        assert 0.45 < metric < 0.55  # Should be close to 0.5
        assert se <= 0.02
        
    def test_adaptive_sampler_small_dataset(self):
        """Test adaptive sampler with small dataset."""
        data = pd.Series(['A', 'B', 'C'] * 10)  # 30 items
        sampler = AdaptiveSampler()
        
        def count_metric(samples):
            return len(set(samples)) / len(samples)
        
        samples, metric, se, info = sampler.sample_until_stable(data, count_metric)
        
        assert info['reason'] == 'small_dataset'
        assert len(samples) == 30  # Should use all data


class TestEarlyAcceptChecker:
    def test_early_accept_clean_data(self):
        """Test early accept with very clean data."""
        metrics = {
            'overlap_ratio': 0.95,
            'variant_rate': 0.02,
            'null_rate': 0.005
        }
        
        should_accept, reason = EarlyAcceptChecker.should_early_accept(metrics, 200)
        assert should_accept is True
        assert reason == "very_clean_data"
        
    def test_early_accept_low_overlap(self):
        """Test early accept with extremely low overlap."""
        metrics = {
            'overlap_ratio': 0.005,
            'variant_rate': 0.1,
            'null_rate': 0.1
        }
        
        should_accept, reason = EarlyAcceptChecker.should_early_accept(metrics, 200)
        assert should_accept is True
        assert reason == "extremely_low_overlap"
        
    def test_no_early_accept(self):
        """Test no early accept for normal data."""
        metrics = {
            'overlap_ratio': 0.5,
            'variant_rate': 0.1,
            'null_rate': 0.05
        }
        
        should_accept, reason = EarlyAcceptChecker.should_early_accept(metrics, 200)
        assert should_accept is False
        assert reason == "continue_sampling"


class TestStratifiedDomainSample:
    def test_stratified_sampling(self):
        """Test stratified sampling maintains TLD distribution."""
        # Create domains with different TLDs
        domains = (
            ['example.com'] * 500 +
            ['test.org'] * 300 +
            ['sample.net'] * 150 +
            ['data.io'] * 50
        )
        series = pd.Series(domains)
        
        sample = stratified_domain_sample(series, target_size=100)
        
        # Check that all TLDs are represented
        sample_tlds = set(d.split('.')[-1] for d in sample)
        assert sample_tlds == {'com', 'org', 'net', 'io'}
        
        # Check approximate proportions
        com_count = sum(1 for d in sample if d.endswith('.com'))
        assert 40 < com_count < 60  # ~50% of sample


class TestDomainWeightCalculation:
    def test_extreme_low_overlap(self):
        """Test weight calculation for extreme low overlap."""
        weight, reason = calculate_domain_weight(
            overlap_ratio=0.005, 
            variant_rate=0.1
        )
        assert weight == 0.1
        assert reason == "extreme_low_overlap"
        
    def test_high_quality_match(self):
        """Test weight calculation for high quality match."""
        weight, reason = calculate_domain_weight(
            overlap_ratio=0.8, 
            variant_rate=0.05
        )
        assert weight == 1.0
        assert reason == "high_quality_match"
        
    def test_normal_fuzzy_match(self):
        """Test weight calculation for normal fuzzy match."""
        weight, reason = calculate_domain_weight(
            overlap_ratio=0.4, 
            variant_rate=0.2
        )
        # Weight should be 0.25 + 0.75 * 0.4 = 0.55
        assert 0.54 <= weight <= 0.56
        assert reason == "fuzzy_match"


class TestDomainAnalysisIntegration:
    def test_analyze_with_adaptive_sampling(self):
        """Test full domain analysis with adaptive sampling."""
        # Create test data
        src = pd.Series(['example.com', 'test.com', 'sample.org'] * 100)
        ref = pd.Series(['example.com', 'test.com', 'different.net'] * 100)
        
        result = analyze_domain_matching_strategy(
            src, ref, use_adaptive=True
        )
        
        assert 'sampling_info' in result['metrics']
        assert result['metrics']['sampling_info']['method'] == 'adaptive'
        assert result['algorithm'] in ['Exact', 'JaroWinkler']
        
    def test_analyze_with_early_accept(self):
        """Test domain analysis with early accept optimization."""
        # Create very clean matching data
        domains = ['example.com', 'test.org', 'sample.net'] * 100
        src = pd.Series(domains)
        ref = pd.Series(domains)  # Identical
        
        result = analyze_domain_matching_strategy(src, ref)
        
        # Should early accept
        assert result['algorithm'] == 'Exact'
        assert result['confidence'] >= 0.95
        assert result['metrics']['sampling_info']['method'] == 'early_accept'


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