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

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

import hist_utils


def perform_hist_modification(
        img_array: np.ndarray,
        hist_ref: dict,
        mode: str,
        L: int = 256
) -> np.ndarray:
    """
    Perform histogram modification on an image according to the desired histogram and selected mode.

    Args:
        img_array (np.ndarray): Input grayscale image with float values in [0, 1].
        hist_ref (dict): Desired output histogram (levels -> relative frequencies).
        mode (str): One of ["greedy", "non-greedy", "post-disturbance"].
        L (int):    Number of levels to clip
    Returns:
        np.ndarray: Modified image array.
    """
    if mode not in ["greedy", "non-greedy", "post-disturbance"]:
        raise ValueError(f"Unknown mode: {mode}")

    # Step 1: Optional post-disturbance (add noise)
    if mode == "post-disturbance":
        unique_levels = np.unique(img_array)
        if len(unique_levels) < 2:
            raise ValueError("Post-disturbance mode requires at least two unique input levels.")

        d = unique_levels[1] - unique_levels[0]
        noise = np.random.uniform(-d / 2, d / 2, size=img_array.shape)

        # Add noise
        img_array = img_array + noise

        # Clip to [0, 1] and quantize to 256 levels to avoid infinite precision
        img_array = np.clip(img_array, 0, 1)
        img_array = np.round(img_array * L) / float(L)

    # Step 2: Calculate input histogram (not normalized)
    input_hist = hist_utils.calculate_hist_of_img(img_array, return_normalized=False)

    # Step 3: Build the modification transform
    modification_transform = {}

    # Prepare input levels (sorted)
    input_levels = sorted(input_hist.keys())
    input_counts = [input_hist[level] for level in input_levels]

    # Prepare target output levels and target counts
    output_levels = sorted(hist_ref.keys())
    output_frequencies = [hist_ref[level] for level in output_levels]
    total_samples = sum(input_counts)
    output_target_counts = [freq * total_samples for freq in output_frequencies]

    # Initialize pointers
    in_idx = 0
    out_idx = 0
    current_in_count = 0
    current_out_target = output_target_counts[out_idx]

    while in_idx < len(input_levels):
        in_level = input_levels[in_idx]
        in_count = input_counts[in_idx]

        if mode in ["greedy", "post-disturbance"]:
            if current_in_count + in_count <= current_out_target:
                modification_transform[in_level] = output_levels[out_idx]
                current_in_count += in_count
                in_idx += 1
            else:
                out_idx += 1
                if out_idx < len(output_levels):
                    current_out_target = output_target_counts[out_idx]
                    current_in_count = 0
                else:
                    # No more output levels: assign the rest to the last output level
                    for remaining_idx in range(in_idx, len(input_levels)):
                        modification_transform[input_levels[remaining_idx]] = output_levels[-1]
                    break

        elif mode == "non-greedy":
            deficiency = current_out_target - current_in_count
            if deficiency >= in_count / 2:
                modification_transform[in_level] = output_levels[out_idx]
                current_in_count += in_count
                in_idx += 1
            else:
                out_idx += 1
                if out_idx < len(output_levels):
                    current_out_target = output_target_counts[out_idx]
                    current_in_count = 0
                else:
                    # No more output levels: assign the rest to the last output level
                    for remaining_idx in range(in_idx, len(input_levels)):
                        modification_transform[input_levels[remaining_idx]] = output_levels[-1]
                    break

    # Step 4: Apply the transform
    modified_img = hist_utils.apply_hist_modification_transform(img_array, modification_transform)

    return modified_img


def perform_hist_eq(
    img_array: np.ndarray,
    mode: str,
    L: int = 256  # number of levels to use for output
) -> np.ndarray:
    """
    Perform histogram equalization on an input image according to the selected mode.

    Args:
        img_array (np.ndarray): Input grayscale image with float values in [0, 1].
        mode (str):             One of ["greedy", "non-greedy", "post-disturbance"].
        L (int):                Number of levels to use for output
    Returns:
        np.ndarray: Equalized image array.
    """
    # Step 1: define L uniformly spaced output levels
    output_levels = np.linspace(0, 1, L)
    hist_ref = {level: 1.0 / L for level in output_levels}

    # Step 2: call the generic modification function
    equalized_img = perform_hist_modification(img_array, hist_ref, mode)

    return equalized_img


def perform_hist_matching(
    img_array: np.ndarray,
    img_array_ref: np.ndarray,
    mode: str
) -> np.ndarray:
    """
    Perform histogram matching of img_array to match img_array_ref according to the selected mode.

    Args:
        img_array (np.ndarray):     Input grayscale image.
        img_array_ref (np.ndarray): Reference grayscale image.
        mode (str):                 One of ["greedy", "non-greedy", "post-disturbance"].

    Returns:
        np.ndarray: Processed image with histogram matched to the reference.
    """
    # Step 1: Calculate the normalized histogram of the reference image
    ref_hist = hist_utils.calculate_hist_of_img(img_array_ref, return_normalized=True)

    # Step 2: Perform histogram modification
    processed_img = perform_hist_modification(img_array, ref_hist, mode)

    return processed_img



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

def test_perform_hist_modification_greedy():
    """
    Test perform_hist_modification with mode 'greedy' on a simple 3x3 image.
    """
    img = np.array([
        [0.0, 0.0, 0.5],
        [0.5, 1.0, 1.0],
        [1.0, 0.5, 0.0]
    ])

    hist_ref = {
        0.25: 0.5,
        0.75: 0.5
    }

    modified_img = perform_hist_modification(img, hist_ref, mode="greedy")

    modified_hist = hist_utils.calculate_hist_of_img(modified_img, return_normalized=True)

    print("Mode: greedy")
    print("Modified image:\n", modified_img)
    print("Modified histogram:", modified_hist)

    output_levels = set(modified_hist.keys())
    expected_levels = {0.25, 0.75}
    print("Test passed:", output_levels == expected_levels)
    print()


def test_perform_hist_modification_non_greedy():
    """
    Test perform_hist_modification with mode 'non-greedy' on a simple 3x3 image.
    """
    img = np.array([
        [0.0, 0.0, 0.5],
        [0.5, 1.0, 1.0],
        [1.0, 0.5, 0.0]
    ])

    hist_ref = {
        0.2: 0.4,
        0.8: 0.6
    }

    modified_img = perform_hist_modification(img, hist_ref, mode="non-greedy")

    modified_hist = hist_utils.calculate_hist_of_img(modified_img, return_normalized=True)

    print("Mode: non-greedy")
    print("Modified image:\n", modified_img)
    print("Modified histogram:", modified_hist)

    output_levels = set(modified_hist.keys())
    expected_levels = {0.2, 0.8}
    print("Test passed:", output_levels == expected_levels)
    print()


def test_perform_hist_modification_post_disturbance():
    """
    Test perform_hist_modification with mode 'post-disturbance' on a simple 3x3 image.
    """
    img = np.array([
        [0.0, 0.0, 0.5],
        [0.5, 1.0, 1.0],
        [1.0, 0.5, 0.0]
    ])

    hist_ref = {
        0.3: 0.5,
        0.7: 0.5
    }

    modified_img = perform_hist_modification(img, hist_ref, mode="post-disturbance")

    modified_hist = hist_utils.calculate_hist_of_img(modified_img, return_normalized=True)

    print("Mode: post-disturbance")
    print("Modified image:\n", modified_img)
    print("Modified histogram:", modified_hist)

    output_levels = set(modified_hist.keys())
    expected_levels = {0.3, 0.7}
    print("Test passed:", output_levels == expected_levels)
    print()



def test_perform_hist_eq_all_modes():
    """
    Test perform_hist_eq for all modes on a simple 3x3 image.
    """
    img = np.array([
        [0.0, 0.0, 0.5],
        [0.5, 1.0, 1.0],
        [1.0, 0.5, 0.0]
    ])

    for mode in ["greedy", "non-greedy", "post-disturbance"]:
        print(f"Mode: {mode} (histogram equalization)")

        equalized_img = perform_hist_eq(img, mode=mode)
        equalized_hist = hist_utils.calculate_hist_of_img(equalized_img, return_normalized=True)

        print("Equalized image:\n", equalized_img)
        print("Equalized histogram:", equalized_hist)

        levels = list(equalized_hist.keys())
        freqs = list(equalized_hist.values())
        expected_freq = 1.0 / len(levels)

        uniform_distribution = all(np.isclose(freq, expected_freq, atol=1e-2) for freq in freqs)
        print("Test passed:", uniform_distribution)
        print()



def test_perform_hist_matching_all_modes():
    """
    Test perform_hist_matching for all modes on a simple 3x3 image.
    """
    img = np.array([
        [0.0, 0.0, 0.5],
        [0.5, 1.0, 1.0],
        [1.0, 0.5, 0.0]
    ])

    # Reference image (different histogram)
    ref_img = np.array([
        [0.2, 0.2, 0.2],
        [0.8, 0.8, 0.8],
        [0.8, 0.2, 0.2]
    ])

    for mode in ["greedy", "non-greedy", "post-disturbance"]:
        print(f"Mode: {mode} (histogram matching)")

        matched_img = perform_hist_matching(img, ref_img, mode=mode)
        matched_hist = hist_utils.calculate_hist_of_img(matched_img, return_normalized=True)
        ref_hist = hist_utils.calculate_hist_of_img(ref_img, return_normalized=True)

        print("Matched image:\n", matched_img)
        print("Matched histogram:", matched_hist)
        print("Reference histogram:", ref_hist)

        matched_levels = set(matched_hist.keys())
        ref_levels = set(ref_hist.keys())

        # Check if output levels are the same as the reference levels
        print("Test passed:", matched_levels == ref_levels)
        print()


def test_perform_hist_matching_simple():
    """
    Simple test for perform_hist_matching function.
    """
    # Create a small input image
    img = np.array([
        [0.0, 0.5],
        [1.0, 0.5]
    ])

    # Create a reference image with known distribution
    ref_img = np.array([
        [0.25, 0.25],
        [0.75, 0.75]
    ])

    # Perform histogram matching using greedy mode
    matched_img = perform_hist_matching(img, ref_img, mode="greedy")

    # Calculate histograms
    matched_hist = hist_utils.calculate_hist_of_img(matched_img, return_normalized=True)
    ref_hist = hist_utils.calculate_hist_of_img(ref_img, return_normalized=True)

    print("Simple test for perform_hist_matching")
    print("Matched image:\n", matched_img)
    print("Matched histogram:", matched_hist)
    print("Reference histogram:", ref_hist)

    # Check if output levels are same as reference levels
    matched_levels = set(matched_hist.keys())
    ref_levels = set(ref_hist.keys())

    print("Test passed:", matched_levels == ref_levels)
    print()


if __name__ == '__main__':
    test_perform_hist_modification_greedy()
    test_perform_hist_modification_non_greedy()
    test_perform_hist_modification_post_disturbance()
    test_perform_hist_eq_all_modes()
    test_perform_hist_matching_all_modes()
    test_perform_hist_matching_simple()