"""
Comprehensive test suite for new distribution data formats in InverseTransformSampler
Tests PMF and histogram formats, plus backward compatibility
"""

import random
from collections import Counter

from samplers import HistogramBucket, InverseTransformSampler, TopNItem


def validate_pmf_format_basic():
    """Test basic PMF (Probability Mass Function) format"""
    print("=== Test PMF Format Basic ===")

    # Simple PMF: dice roll
    pmf_data = {
        "type": "pmf",
        "values": [1, 2, 3, 4, 5, 6],
        "probabilities": [1 / 6, 1 / 6, 1 / 6, 1 / 6, 1 / 6, 1 / 6],
    }

    sampler = InverseTransformSampler(distribution_data=pmf_data, seed=42)

    # Test sampling
    sample_results = sampler.sample_batch(6000)
    non_null_samples = [s.value for s in sample_results if s.value is not None]

    print(
        f"  Generated {len(sample_results)} samples, {len(non_null_samples)} non-null"
    )

    # Check value distribution
    value_counts = Counter(non_null_samples)
    print(f"  Value distribution: {dict(value_counts)}")

    # Each value should appear roughly 1000 times (6000 / 6)
    for value in [1, 2, 3, 4, 5, 6]:
        count = value_counts.get(value, 0)
        expected = 1000
        error = abs(count - expected) / expected
        print(f"    Value {value}: {count} occurrences (error: {error:.1%})")
        assert error < 0.1, f"PMF distribution error too high for value {value}"

    print("  ✅ PMF basic test passed")
    return sampler


def validate_pmf_format_weighted():
    """Test PMF with non-uniform probabilities"""
    print("\n=== Test PMF Format Weighted ===")

    # Weighted PMF: more likely to get low numbers
    pmf_data = {
        "type": "pmf",
        "values": [1, 2, 3, 4, 5],
        "probabilities": [0.5, 0.25, 0.15, 0.07, 0.03],
    }

    sampler = InverseTransformSampler(distribution_data=pmf_data, seed=42)

    # Test sampling
    sample_results = sampler.sample_batch(10000)
    non_null_samples = [s.value for s in sample_results if s.value is not None]

    print(
        f"  Generated {len(sample_results)} samples, {len(non_null_samples)} non-null"
    )

    # Check value distribution matches probabilities
    value_counts = Counter(non_null_samples)
    total_samples = len(non_null_samples)

    expected_probs = {1: 0.5, 2: 0.25, 3: 0.15, 4: 0.07, 5: 0.03}

    for value, expected_prob in expected_probs.items():
        count = value_counts.get(value, 0)
        actual_prob = count / total_samples
        error = abs(actual_prob - expected_prob) / expected_prob
        print(
            f"    Value {value}: {actual_prob:.3f} prob (expected {expected_prob:.3f}, error: {error:.1%})"
        )
        assert error < 0.1, f"PMF probability error too high for value {value}"

    print("  ✅ PMF weighted test passed")
    return sampler


def validate_histogram_format_basic():
    """Test basic histogram format"""
    print("\n=== Test Histogram Format Basic ===")

    # Histogram: age distribution
    hist_data = {
        "type": "histogram",
        "bins": [0, 20, 40, 60, 80, 100],  # Age bins
        "counts": [100, 300, 400, 150, 50],  # Counts per bin
    }

    sampler = InverseTransformSampler(distribution_data=hist_data, seed=42)

    # Test sampling
    sample_results = sampler.sample_batch(5000)
    samples = [s.value for s in sample_results]
    non_null_samples = [s for s in samples if s is not None]

    print(f"  Generated {len(samples)} samples, {len(non_null_samples)} non-null")

    # Check that samples fall within expected ranges
    bins = [(0, 20), (20, 40), (40, 60), (60, 80), (80, 100)]
    bin_counts = [0] * len(bins)

    for sample in non_null_samples:
        for i, (lower, upper) in enumerate(bins):
            if lower <= sample < upper or (i == len(bins) - 1 and sample == upper):
                bin_counts[i] += 1
                break

    print(f"  Bin distribution: {bin_counts}")

    # Check relative distribution matches expected
    expected_counts = [100, 300, 400, 150, 50]
    total_expected = sum(expected_counts)

    for i, (expected, actual) in enumerate(zip(expected_counts, bin_counts)):
        expected_ratio = expected / total_expected
        actual_ratio = actual / len(non_null_samples)
        error = (
            abs(actual_ratio - expected_ratio) / expected_ratio
            if expected_ratio > 0
            else 0
        )
        print(
            f"    Bin {bins[i]}: {actual_ratio:.3f} ratio (expected {expected_ratio:.3f}, error: {error:.1%})"
        )
        if expected > 0:  # Only check non-empty bins
            assert error < 0.15, f"Histogram bin error too high for bin {i}"

    print("  ✅ Histogram basic test passed")
    return sampler


def validate_histogram_format_continuous():
    """Test histogram with continuous-like distribution"""
    print("\n=== Test Histogram Format Continuous ===")

    # Normal-like distribution histogram
    hist_data = {
        "type": "histogram",
        "bins": [-3, -2, -1, 0, 1, 2, 3],
        "counts": [50, 150, 300, 300, 150, 50],  # Bell curve shape
    }

    sampler = InverseTransformSampler(
        distribution_data=hist_data, null_count=100, seed=42
    )

    # Test sampling
    sample_results = sampler.sample_batch(5000)
    samples = [s.value for s in sample_results]
    non_null_samples = [s for s in samples if s is not None]
    null_count = len(samples) - len(non_null_samples)

    print(
        f"  Generated {len(samples)} samples, {len(non_null_samples)} non-null, {null_count} null"
    )

    # Check range
    if non_null_samples:
        min_sample = min(non_null_samples)
        max_sample = max(non_null_samples)
        print(f"  Sample range: {min_sample:.3f} to {max_sample:.3f}")

        # Should be roughly within [-3, 3]
        assert (
            -3.1 <= min_sample <= -2.9
        ), f"Min sample out of expected range: {min_sample}"
        assert (
            2.9 <= max_sample <= 3.1
        ), f"Max sample out of expected range: {max_sample}"

    # Check null probability
    expected_null_rate = 100 / (1000 + 100)  # null_count / (sum(counts) + null_count)
    actual_null_rate = null_count / len(samples)
    error = abs(actual_null_rate - expected_null_rate) / expected_null_rate
    print(
        f"  Null rate: {actual_null_rate:.3f} (expected {expected_null_rate:.3f}, error: {error:.1%})"
    )
    assert error < 0.1, f"Null rate error too high: {error}"

    print("  ✅ Histogram continuous test passed")
    return sampler


def validate_backward_compatibility():
    """Test that old buckets/top_n format still works"""
    print("\n=== Test Backward Compatibility ===")

    # Original format
    buckets = [
        HistogramBucket(0.0, 10.0, 100, 10),
        HistogramBucket(10.0, 20.0, 200, 20),
    ]

    top_n = [TopNItem(5.0, 80), TopNItem(15.0, 60)]

    # Test old constructor style
    old_sampler = InverseTransformSampler(
        buckets=buckets, top_n=top_n, null_count=20, is_cumulative_count=False, seed=42
    )
    old_sample_results = old_sampler.sample_batch(1000)
    old_non_null = [s.value for s in old_sample_results if s.value is not None]

    print(
        f"  Old format: {len(old_sample_results)} samples, {len(old_non_null)} non-null"
    )

    # Test new constructor with None distribution_data (should use buckets/top_n)
    new_sampler = InverseTransformSampler(
        distribution_data=None,
        buckets=buckets,
        top_n=top_n,
        null_count=20,
        is_cumulative_count=False,
        seed=42,
    )
    new_sample_results = new_sampler.sample_batch(1000)
    new_non_null = [s.value for s in new_sample_results if s.value is not None]

    print(
        f"  New format (explicit None): {len(new_sample_results)} samples, {len(new_non_null)} non-null"
    )

    # Should produce similar distributions
    old_counts = Counter(old_non_null)
    new_counts = Counter(new_non_null)

    print(
        f"  Old TopN frequencies: 5.0={old_counts.get(5.0, 0)}, 15.0={old_counts.get(15.0, 0)}"
    )
    print(
        f"  New TopN frequencies: 5.0={new_counts.get(5.0, 0)}, 15.0={new_counts.get(15.0, 0)}"
    )

    print("  ✅ Backward compatibility test passed")


def validate_error_handling():
    """Test error handling for invalid distribution data"""
    print("\n=== Test Error Handling ===")

    # Test invalid format
    try:
        InverseTransformSampler(distribution_data={"invalid": "format"})
        assert False, "Should have raised ValueError for missing 'type'"
    except ValueError as e:
        print(f"  ✅ Caught expected error for missing 'type': {e}")

    # Test unsupported type
    try:
        InverseTransformSampler(distribution_data={"type": "unsupported"})
        assert False, "Should have raised ValueError for unsupported type"
    except ValueError as e:
        print(f"  ✅ Caught expected error for unsupported type: {e}")

    # Test PMF with mismatched lengths
    try:
        InverseTransformSampler(
            distribution_data={
                "type": "pmf",
                "values": [1, 2, 3],
                "probabilities": [0.5, 0.5],  # Wrong length
            }
        )
        assert False, "Should have raised ValueError for mismatched lengths"
    except ValueError as e:
        print(f"  ✅ Caught expected error for mismatched lengths: {e}")

    # Test PMF with probabilities not summing to 1
    try:
        InverseTransformSampler(
            distribution_data={
                "type": "pmf",
                "values": [1, 2],
                "probabilities": [0.5, 0.6],  # Sum = 1.1
            }
        )
        assert False, "Should have raised ValueError for probabilities not summing to 1"
    except ValueError as e:
        print(f"  ✅ Caught expected error for invalid probability sum: {e}")

    # Test histogram with wrong bin/count relationship
    try:
        InverseTransformSampler(
            distribution_data={
                "type": "histogram",
                "bins": [0, 10, 20],  # 3 bins
                "counts": [100, 200, 300],  # 3 counts (should be 2)
            }
        )
        assert False, "Should have raised ValueError for wrong bin/count relationship"
    except ValueError as e:
        print(f"  ✅ Caught expected error for wrong bin/count relationship: {e}")

    print("  ✅ Error handling test passed")


def validate_edge_cases():
    """Test edge cases and boundary conditions"""
    print("\n=== Test Edge Cases ===")

    # Empty PMF
    empty_pmf = InverseTransformSampler(
        distribution_data={"type": "pmf", "values": [], "probabilities": []},
        null_count=10,
        seed=42,
    )

    empty_sample_results = empty_pmf.sample_batch(10)
    empty_samples = [s.value for s in empty_sample_results]
    print(f"  Empty PMF: {empty_samples}")
    assert all(s is None for s in empty_samples), "Empty PMF should return only nulls"

    # Single value PMF
    single_pmf = InverseTransformSampler(
        distribution_data={"type": "pmf", "values": [42], "probabilities": [1.0]},
        seed=42,
    )

    single_sample_results = single_pmf.sample_batch(10)
    single_samples = [s.value for s in single_sample_results]
    single_non_null = [s for s in single_samples if s is not None]
    print(f"  Single PMF: {single_non_null[:5]}")
    assert all(
        s == 42.0 for s in single_non_null
    ), "Single PMF should return only the single value"

    # Empty histogram (no counts)
    empty_hist = InverseTransformSampler(
        distribution_data={"type": "histogram", "bins": [0, 10], "counts": []},
        null_count=5,
        seed=42,
    )

    empty_hist_sample_results = empty_hist.sample_batch(5)
    empty_hist_samples = [s.value for s in empty_hist_sample_results]
    print(f"  Empty histogram: {empty_hist_samples}")
    assert all(
        s is None for s in empty_hist_samples
    ), "Empty histogram should return only nulls"

    # Zero counts in histogram
    zero_hist = InverseTransformSampler(
        distribution_data={
            "type": "histogram",
            "bins": [0, 10, 20, 30],
            "counts": [0, 100, 0],  # Only middle bin has data
        },
        seed=42,
    )

    zero_sample_results = zero_hist.sample_batch(100)
    zero_samples = [s.value for s in zero_sample_results]
    zero_non_null = [s for s in zero_samples if s is not None]
    print(f"  Zero count bins: {len(zero_non_null)} non-null samples")
    if zero_non_null:
        print(f"    Range: {min(zero_non_null):.3f} to {max(zero_non_null):.3f}")
        # All should be in [10, 20) range
        assert all(
            10 <= s < 20 for s in zero_non_null
        ), "Should only sample from non-zero bin"

    print("  ✅ Edge cases test passed")


def validate_performance_comparison():
    """Compare performance of different formats"""
    print("\n=== Test Performance Comparison ===")

    import time

    # Create equivalent distributions in different formats
    # PMF format
    pmf_data = {
        "type": "pmf",
        "values": list(range(1, 101)),  # 1 to 100
        "probabilities": [0.01] * 100,  # Uniform
    }

    # Histogram format
    hist_data = {
        "type": "histogram",
        "bins": list(range(0, 102)),  # 0 to 101
        "counts": [100] * 101,  # Uniform
    }

    # Old format
    buckets = [HistogramBucket(float(i), float(i + 1), 100, 0) for i in range(100)]

    samplers = [
        ("PMF", InverseTransformSampler(distribution_data=pmf_data, seed=42)),
        ("Histogram", InverseTransformSampler(distribution_data=hist_data, seed=42)),
        ("TiDB", InverseTransformSampler(buckets=buckets, seed=42)),
    ]

    num_samples = 50000

    for name, sampler in samplers:
        start_time = time.time()
        sample_results = sampler.sample_batch(num_samples)
        samples = [s.value for s in sample_results]
        end_time = time.time()

        elapsed = end_time - start_time
        rate = num_samples / elapsed
        non_null = sum(1 for s in samples if s is not None)

        print(
            f"  {name} format: {elapsed:.4f}s, {rate:,.0f} samples/sec, {non_null} non-null"
        )

    print("  ✅ Performance comparison completed")


def run_comprehensive_distribution_format_test():
    """Run comprehensive test suite for new distribution formats"""
    print("🚀 Starting Comprehensive Distribution Format Test Suite")
    print("=" * 80)

    # Set random seed for reproducibility
    random.seed(42)

    try:
        # 1. PMF format tests
        validate_pmf_format_basic()
        validate_pmf_format_weighted()

        # 2. Histogram format tests
        validate_histogram_format_basic()
        validate_histogram_format_continuous()

        # 3. Backward compatibility
        validate_backward_compatibility()

        # 4. Error handling
        validate_error_handling()

        # 5. Edge cases
        validate_edge_cases()

        # 6. Performance comparison
        validate_performance_comparison()

        print("\n" + "=" * 80)
        print("🎉 Comprehensive Distribution Format Test Suite Completed!")

        print("\n🎯 Test Results Summary:")
        print("  ✅ PMF format: Supports exact probability specification")
        print("  ✅ Histogram format: Supports binned continuous distributions")
        print("  ✅ Backward compatibility: Old TiDB format still works")
        print("  ✅ Error handling: Proper validation and error messages")
        print("  ✅ Edge cases: Handles empty and boundary conditions")
        print("  ✅ Performance: All formats maintain high sampling speed")

        print("\n💡 Format Recommendations:")
        print("  • Use PMF for exact discrete values with known probabilities")
        print("  • Use Histogram for continuous/binned data")
        print("  • Use TiDB format for database statistics (backward compatible)")
        print("  • All formats support null values and maintain statistical accuracy")

        return True

    except Exception as e:
        print(f"\n❌ Distribution format test failed: {e}")
        import traceback

        traceback.print_exc()
        return False


if __name__ == "__main__":
    success = run_comprehensive_distribution_format_test()
    exit(0 if success else 1)