import numpy as np
from granatpy.core import naturalize_rgb_image

def test_naturalize_rgb_image():
    # Create a random RGB image (10x10)
    image = np.random.randint(0, 256, (10, 10, 3), dtype=np.uint8)
    
    result, nfs, t1s, t2s, psnrs = naturalize_rgb_image(image)
    
    # Assert return types and shapes
    assert isinstance(result, np.ndarray)
    assert result.shape == (10, 10, 3)
    assert result.dtype == np.uint8
    
    assert isinstance(nfs, list) and len(nfs) == 3
    assert isinstance(t1s, list) and len(t1s) == 3
    assert isinstance(t2s, list) and len(t2s) == 3
    assert isinstance(psnrs, list) and len(psnrs) == 3


def test_naturalize_grayscale_image():
    # Create a random grayscale image (10x10)
    image = np.random.randint(0, 256, (10, 10), dtype=np.uint8)
    
    result, nfs, t1s, t2s, psnrs = naturalize_rgb_image(image)
    
    # Assert return types and shapes
    assert isinstance(result, np.ndarray)
    assert result.shape == (10, 10)
    assert result.dtype == np.uint8
    
    assert isinstance(nfs, list) and len(nfs) == 1
    assert isinstance(t1s, list) and len(t1s) == 1
    assert isinstance(t2s, list) and len(t2s) == 1
    assert isinstance(psnrs, list) and len(psnrs) == 1
