#
# histogram utilities file
#
#
# author: Christos Choutouridis <cchoutou@ece.auth.gr>
# date:   29/04/2025
#

try:
    import numpy as np
    from PIL import Image
except ImportError as e:
    print("Missing package: ", e)
    print("Run: pip install -r requirements.txt to install.")
    exit(1)


def calculate_hist_of_img(
        img_array: np.ndarray,
        return_normalized: bool
) -> dict:
    """
    Calculate the histogram of an input image array.

    Args:
        img_array (np.ndarray):   2D grayscale image array with float values in [0, 1].
        return_normalized (bool): If True, return normalized histogram (probabilities),
                                  otherwise return absolute counts.

    Returns:
        dict: A dictionary mapping each unique input level to its count or probability.
    """
    # Flatten the 2D image into a 1D array for easier processing
    flat_img = img_array.flatten()

    # Create an empty dictionary to store the histogram
    hist = {}

    # Traverse each pixel value
    for value in flat_img:
        # Increment the count for the corresponding pixel value
        if value in hist:
            hist[value] += 1
        else:
            hist[value] = 1

    if return_normalized:
        total_pixels = flat_img.size
        for key in hist:
            hist[key] /= total_pixels

    return hist



def apply_hist_modification_transform(
        img_array: np.ndarray,
        modification_transform: dict
) -> np.ndarray:
    """
    Apply a histogram modification transform to an image.

    Args:
        img_array (np.ndarray):        2D grayscale image array with float values.
        modification_transform (dict): Dictionary mapping each input level to an output level.

    Returns:
        np.ndarray: The modified image array.
    """
    # Create a copy of the input image to avoid modifying it in-place
    modified_img = np.copy(img_array)

    # Get the unique levels in the image
    unique_levels = np.unique(modified_img)

    # Apply the transform
    for level in unique_levels:
        if level in modification_transform:
            modified_img[modified_img == level] = modification_transform[level]
        else:
            # Optional: raise an error or warning if level is missing
            raise ValueError(f"Input level {level} not found in modification_transform.")

    return modified_img









#
# =========================== Functional tests ===========================
#
# To execute the tests run:
#   python hist_utils.py
#

def test_calculate_hist_counts():
    """
    Test calculate_hist_of_img with return_normalized = False.
    """
    # Create a dummy 3x3 image
    img = np.array([
        [0.0, 0.5, 1.0],
        [0.5, 0.5, 0.0],
        [1.0, 0.0, 1.0]
    ])

    # Call our function
    my_hist = calculate_hist_of_img(img, return_normalized=False)

    # Use numpy to verify
    values, counts = np.unique(img, return_counts=True)
    np_hist = dict(zip(values, counts))

    # Compare
    print("Test case: Counts")
    #print("Our histogram:", my_hist)
    #print("Numpy histogram:", np_hist)
    print("Test passed:", my_hist == np_hist)
    print()


def test_calculate_hist_normalized():
    """
    Test calculate_hist_of_img with return_normalized = True.
    """
    # Create a dummy 3x3 image
    img = np.array([
        [0.0, 0.5, 1.0],
        [0.5, 0.5, 0.0],
        [1.0, 0.0, 1.0]
    ])

    # Call our function
    my_hist = calculate_hist_of_img(img, return_normalized=True)

    # Use numpy to verify (manual normalization)
    values, counts = np.unique(img, return_counts=True)
    total_pixels = img.size
    np_hist = dict(zip(values, counts / total_pixels))

    # Compare
    print("Test case: Normalized")
    #print("Our histogram:", my_hist)
    #print("Numpy histogram:", np_hist)

    # Because of floating point division, use np.isclose instead of exact equality
    test_passed = True
    for key in my_hist:
        if not np.isclose(my_hist[key], np_hist[key]):
            test_passed = False
            break

    print("Test passed:", test_passed)
    print()



def test_apply_hist_modification_identity():
    """
    Test apply_hist_modification_transform with identity transform (no changes).
    """
    # Create a simple 3x3 image
    img = np.array([
        [0.0, 0.5, 1.0],
        [0.5, 0.5, 0.0],
        [1.0, 0.0, 1.0]
    ])

    # Define an identity transform (each value maps to itself)
    mod_transform = {
        0.0: 0.0,
        0.5: 0.5,
        1.0: 1.0
    }

    # Apply our function
    modified_img = apply_hist_modification_transform(img, mod_transform)

    # Compare
    print("Test case: Identity transform")
    # print("Modified image:\n", modified_img)
    # print("Original image:\n", img)
    print("Test passed:", np.allclose(modified_img, img))
    print()


def test_apply_hist_modification_simple():
    """
    Test apply_hist_modification_transform with a simple transform.
    """
    # Create a simple 3x3 image
    img = np.array([
        [0.0, 0.5, 1.0],
        [0.5, 0.5, 0.0],
        [1.0, 0.0, 1.0]
    ])

    # Define a simple modification transform
    mod_transform = {
        0.0: 0.2,
        0.5: 0.7,
        1.0: 1.0
    }

    # Expected result
    expected_img = np.array([
        [0.2, 0.7, 1.0],
        [0.7, 0.7, 0.2],
        [1.0, 0.2, 1.0]
    ])

    # Apply our function
    modified_img = apply_hist_modification_transform(img, mod_transform)

    # Compare
    print("Test case: Simple transform")
    # print("Modified image:\n", modified_img)
    # print("Expected image:\n", expected_img)
    print("Test passed:", np.allclose(modified_img, expected_img))
    print()


if __name__ == '__main__':
    test_calculate_hist_counts()
    test_calculate_hist_normalized()
    test_apply_hist_modification_identity()
    test_apply_hist_modification_simple()