#
# Image to graph utility
#
# For the given data we have:
# dip_hw_3.mat["d1a"]    # [MN x MN] affinity matrix
# dip_hw_3.mat["d2a"]    # [M x N x 3] RGB image
# dip_hw_3.mat["d2b"]    # [M x N x 3] RGB image
#
#
# author: Christos Choutouridis <cchoutou@ece.auth.gr>
# date:   05/07/2025
#

try:
    import numpy as np
    from numpy._typing import NDArray
    from sklearn.metrics import pairwise_distances

    # Testing requirements
    import matplotlib.pyplot
    from scipy.io import loadmat
except ImportError as e:
    print("Missing package: ", e)
    print("Run: pip install -r requirements.txt to install.")
    exit(1)

def image_to_graph(
    img_array: NDArray[np.floating]
) -> NDArray[np.float64]:
    """
    Converts an input image into a fully connected graph represented
    by an affinity matrix.

    Parameters:
    ----------
    img_array : np.ndarray of shape (M, N, C), dtype=float
        The input image with C channels (e.g., 3 for RGB),
        with values normalized in [0, 1].

    Returns:
    -------
    affinity_mat : np.ndarray of shape (M*N, M*N), dtype=float
        Symmetric affinity matrix representing the fully connected graph.
        A(i, j) = 1 / ||pixel_i - pixel_j||_2
    """
    if not np.issubdtype(img_array.dtype, np.floating):
        raise ValueError("img_array must be of float type with values in [0, 1].")

    M, N, C = img_array.shape
    pixels = img_array.reshape(-1, C)  # shape (M*N, C)

    # Compute Euclidean distances between all pixel vectors
    distances = pairwise_distances(pixels, metric='euclidean')  # shape (MN, MN)

    # Avoid division by zero on the diagonal
    np.fill_diagonal(distances, 1e-10)

    # Affinity = 1 / e^d(i,j)
    affinity_mat = 1.0 / np.exp(distances)

    return affinity_mat



def _test_1(plot : bool = False):
    """
    Test image_to_graph() with a small 4x4 RGB random value array
    """
    print(f" === Test 1 === ")
    print(f"")
    # Small 4x4 RGB with random values at [0, 1]
    img_array = np.random.rand(4, 4, 3).astype(np.float32)

    # affinity matrix calculation
    A = image_to_graph(img_array)

    # Print specs
    print("Shape of affinity matrix:", A.shape)  # (16, 16)
    print("Is symmetric:", np.allclose(A, A.T))  # True
    print("Max value:", np.max(A))
    print("Min value:", np.min(A))

    if plot:
        matplotlib.use("TkAgg")
        matplotlib.pyplot.imshow(A, cmap='hot')
        matplotlib.pyplot.colorbar()
        matplotlib.pyplot.title("Affinity Matrix Heatmap")
        matplotlib.pyplot.show()


def _test_2(plot : bool = False):
    """
    Test image_to_graph() with d2b matrix
    """
    print(f" === Test 2 === ")
    print(f"")

    data = loadmat("dip_hw_3.mat")
    img_array = data["d2b"]  # shape (M, N, 3), dtype float, values in [0, 1]

    # Check shape and type
    print("Input image shape:", img_array.shape)
    print("dtype:", img_array.dtype)

    # affinity matrix calculation
    A = image_to_graph(img_array)

    # Print specs
    print("Affinity matrix shape:", A.shape)
    print("Is symmetric:", np.allclose(A, A.T))
    print("Max value:", np.max(A))
    print("Min value:", np.min(A))

    if plot:
        matplotlib.use("TkAgg")
        matplotlib.pyplot.imshow(A, cmap='hot')
        matplotlib.pyplot.title("Affinity Matrix Heatmap")
        matplotlib.pyplot.colorbar()
        matplotlib.pyplot.show()


if __name__ == '__main__':
    # If you have TkAgg you can pass True, otherwise pass False
    _test_1(False)
    _test_2(False)