#!/usr/bin/env python3
"""Test the enumerative entropy coding implementation."""

from enumerative_coding import EnumerativeEncoder, ExpGolombCoder


def test_exp_golomb():
    """Test exp-Golomb coding."""
    print("Testing Exp-Golomb coding:")
    
    test_values = [0, 1, 2, 3, 4, 5, 10, 15, 31, 32, 100]
    
    for n in test_values:
        encoded = ExpGolombCoder.encode(n)
        decoded, _ = ExpGolombCoder.decode(encoded, 0)
        print(f"  {n:3d} -> {encoded:>10s} -> {decoded:3d} ({'✓' if n == decoded else '✗'})")


def test_combinatorics():
    """Test combinatorial functions."""
    print("\nTesting combinatorial functions:")
    
    # Test binomial coefficients using the encoder's table
    encoder = EnumerativeEncoder()
    test_cases = [(5, 2), (10, 3), (7, 0), (7, 7), (6, 4)]
    for n, k in test_cases:
        result = encoder.binom_table.get(n, k)
        expected = math.comb(n, k) if hasattr(math, 'comb') else None
        print(f"  C({n},{k}) = {result}" + (f" (expected {expected})" if expected else ""))


def test_simple_sequence():
    """Test encoding/decoding of simple sequences."""
    print("\nTesting simple sequences:")
    
    encoder = EnumerativeEncoder()
    
    test_sequences = [
        [0, 1, 0],
        [0, 1, 1, 2],
        [1, 2, 3, 1, 2, 3],
        [0, 0, 1, 1, 2],
    ]
    
    for seq in test_sequences:
        print(f"\n  Testing sequence: {seq}")
        
        try:
            encoded = encoder.encode(seq)
            decoded = encoder.decode(encoded)
            
            print(f"    Original:  {seq}")
            print(f"    Decoded:   {decoded}")
            print(f"    Correct:   {'✓' if seq == decoded else '✗'}")
            print(f"    Size:      {len(seq)} symbols -> {len(encoded)} bytes")
            
        except Exception as e:
            print(f"    Error: {e}")


def test_paper_example():
    """Test with an example that should match the paper's approach."""
    print("\nTesting paper-style example:")
    
    # Create a sequence with known symbol frequencies
    # This tests the 3-step encoding process
    sequence = [0, 0, 1, 2, 1, 0, 2, 1, 1, 2]  # 3 zeros, 4 ones, 3 twos
    
    print(f"  Sequence: {sequence}")
    print(f"  Symbols:  {sorted(set(sequence))}")
    print(f"  Counts:   {[sequence.count(s) for s in sorted(set(sequence))]}")
    
    encoder = EnumerativeEncoder()
    encoded = encoder.encode(sequence)
    decoded = encoder.decode(encoded)
    
    print(f"  Encoded size: {len(encoded)} bytes")
    print(f"  Correctly decoded: {'✓' if sequence == decoded else '✗'}")
    
    if sequence != decoded:
        print(f"  Expected: {sequence}")
        print(f"  Got:      {decoded}")


if __name__ == "__main__":
    import math
    
    test_exp_golomb()
    test_combinatorics()
    test_simple_sequence()
    test_paper_example()