#
# Demo 3a: Non-recursive normalized cuts
#
# author: Christos Choutouridis <cchoutou@ece.auth.gr>
# date:   06/07/2025
#
try:
    # Testing requirements
    from scipy.io import loadmat
    import matplotlib.pyplot as plt
    import numpy as np

    # Project imports
    from normalized_cuts import n_cuts
    from spectral_clustering import spectral_clustering
    from image_to_graph import image_to_graph
except ImportError as e:
    print("Missing package:", e)
    print("Run: pip install -r requirements.txt")
    exit(1)




def plot_segmentation(image, labels, k, title, fname):
    M, N, _ = image.shape
    segmented = labels.reshape(M, N)
    plt.imshow(segmented, cmap='tab10', vmin=0, vmax=k-1)
    plt.title(title)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig(fname)
    print(f"Saved: {fname}")
    plt.close()


def run_demo3a():
    data = loadmat("dip_hw_3.mat")

    for name in ["d2a", "d2b"]:
        img = data[name]
        print(f"\n=== Image {name} ===")

        affinity = image_to_graph(img)

        for k in [2, 3, 4]:
            print(f"  k = {k}")

            labels_nc = n_cuts(affinity, k=k)
            labels_sc = spectral_clustering(affinity, k=k, normalized=False)
            labels_sc_nrm = spectral_clustering(affinity, k=k, normalized=True)

            plot_segmentation(img, labels_nc, k, f"{name} - n_cuts (k={k})", f"plots/demo3a_{name}_ncuts_k{k}.png")
            plot_segmentation(img, labels_sc, k, f"{name} - spectral (k={k})", f"plots/demo3a_{name}_spectral_k{k}.png")
            plot_segmentation(img, labels_sc_nrm, k, f"{name} - spectral-Lnorm (k={k})", f"plots/demo3a_{name}_spectral_k{k}_norm.png")


if __name__ == '__main__':
    run_demo3a()