"""
Test image compression and decompression roundtrip.
"""

import numpy as np
import pytest
from pathlib import Path

from sempress import (
    encode_image_from_array,
    decode_image,
    ImageEncodeConfig,
    psnr,
    ssim
)


def test_image_roundtrip_basic():
    """Test basic image compression and decompression."""
    # Create a simple test image (100x100 RGB)
    img = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)
    
    # Configure compression
    config = ImageEncodeConfig(
        patch_size=8,
        k=64,  # Small for faster test
        color_space="RGB",
        quality_mode="fast"
    )
    
    # Compress
    compressed = encode_image_from_array(img, config)
    
    # Decompress
    reconstructed = decode_image(compressed)
    
    # Check shape matches
    assert reconstructed.shape == img.shape
    
    # Note: Random noise doesn't compress well, so we just verify roundtrip works
    # The compressed size may be larger due to codebook overhead
    assert len(compressed) > 0
    
    # Check quality (should be reasonable for random noise)
    psnr_val = psnr(img, reconstructed)
    assert psnr_val > 10  # Very loose bound for random data


def test_image_roundtrip_smooth():
    """Test image compression on smooth image (better compression)."""
    # Create a smooth gradient image
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    for i in range(100):
        img[i, :, 0] = i * 2  # Red gradient
        img[:, i, 1] = i * 2  # Green gradient
    img[:, :, 2] = 128  # Constant blue
    
    config = ImageEncodeConfig(
        patch_size=8,
        k=128,
        color_space="RGB",
        quality_mode="balanced"
    )
    
    compressed = encode_image_from_array(img, config)
    reconstructed = decode_image(compressed)
    
    # Should achieve good compression on smooth data
    compression_ratio = img.nbytes / len(compressed)
    assert compression_ratio > 2.0
    
    # Quality should be high on smooth data
    psnr_val = psnr(img, reconstructed)
    ssim_val = ssim(img, reconstructed)
    
    assert psnr_val > 25  # Should be high for smooth gradient
    assert ssim_val > 0.7  # Should preserve structure


def test_image_grayscale():
    """Test grayscale image compression."""
    # Create grayscale image - simple gradient for better compression
    img = np.zeros((100, 100), dtype=np.uint8)
    for i in range(100):
        img[i, :] = i * 2  # Vertical gradient
    
    # Convert to RGB for processing (grayscale mode expects RGB input)
    img_rgb = np.stack([img, img, img], axis=2)
    
    config = ImageEncodeConfig(
        patch_size=8,
        k=64,
        color_space="Grayscale",
        quality_mode="fast"
    )
    
    compressed = encode_image_from_array(img_rgb, config)
    reconstructed = decode_image(compressed)
    
    # Should produce valid image (will be 2D for grayscale)
    assert reconstructed.ndim in [2, 3]  # Can be 2D or 3D
    assert reconstructed.shape[:2] == img.shape
    
    # Should compress smooth gradient well
    assert len(compressed) < img_rgb.nbytes


def test_image_ycbcr():
    """Test YCbCr color space compression."""
    # Create smooth gradient for better compression
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    for i in range(100):
        img[i, :, 0] = i * 2  # Red gradient
        img[:, i, 1] = i * 2  # Green gradient
    img[:, :, 2] = 128  # Constant blue
    
    config = ImageEncodeConfig(
        patch_size=8,
        k=64,
        color_space="YCbCr",
        quality_mode="fast"
    )
    
    compressed = encode_image_from_array(img, config)
    reconstructed = decode_image(compressed)
    
    assert reconstructed.shape == img.shape
    # Smooth gradient should compress well
    assert len(compressed) < img.nbytes


def test_image_different_patch_sizes():
    """Test different patch sizes."""
    # Create smooth gradient for better compression
    img = np.zeros((128, 128, 3), dtype=np.uint8)
    for i in range(128):
        img[i, :, 0] = i * 2  # Red gradient
        img[:, i, 1] = i * 2  # Green gradient
    img[:, :, 2] = 128  # Constant blue
    
    for patch_size in [8, 16]:
        config = ImageEncodeConfig(
            patch_size=patch_size,
            k=64,
            quality_mode="fast"
        )
        
        compressed = encode_image_from_array(img, config)
        reconstructed = decode_image(compressed)
        
        assert reconstructed.shape == img.shape
        # Smooth gradient should compress well
        assert len(compressed) < img.nbytes


def test_image_residuals():
    """Test residual storage for high-error patches."""
    # Create image with some variation (but not pure random)
    img = np.zeros((100, 100, 3), dtype=np.uint8)
    # Create gradient with some random variation
    for i in range(100):
        for j in range(100):
            img[i, j, 0] = (i + np.random.randint(-20, 20)) % 256
            img[i, j, 1] = (j + np.random.randint(-20, 20)) % 256
            img[i, j, 2] = 128
    
    # Low threshold = more residuals stored
    config_with_residuals = ImageEncodeConfig(
        patch_size=8,
        k=32,
        residual_threshold=0.05,  # Store residuals for 5% error
        quality_mode="high"
    )
    
    # High threshold = fewer residuals
    config_without_residuals = ImageEncodeConfig(
        patch_size=8,
        k=32,
        residual_threshold=0.5,  # Store residuals only for 50% error
        quality_mode="fast"
    )
    
    compressed_with = encode_image_from_array(img, config_with_residuals)
    compressed_without = encode_image_from_array(img, config_without_residuals)
    
    # Version with residuals should be larger but better quality
    assert len(compressed_with) > len(compressed_without)
    
    recon_with = decode_image(compressed_with)
    recon_without = decode_image(compressed_without)
    
    psnr_with = psnr(img, recon_with)
    psnr_without = psnr(img, recon_without)
    
    # More residuals = better quality (or at least not worse)
    assert psnr_with >= psnr_without - 1.0  # Allow small margin


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