#
# demo 1 of the assignment
#
# 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:
    # Testing requirements
    import numpy as np
    from scipy.io import loadmat
    import matplotlib.pyplot
    # Project imports
    from spectral_clustering import *
except ImportError as e:
    print("Missing package:", e)
    print("Run: pip install -r requirements.txt")
    exit(1)


def run_demo1():
    data = loadmat("dip_hw_3.mat")
    A = data["d1a"]
    print("Loaded affinity matrix d1a with shape:", A.shape)

    for k in [2, 3, 4]:
        print(f"\n=== Spectral Clustering on d1a with k={k} ===")

        labels = spectral_clustering(A, k)
        print("Cluster labels:")
        print(labels)

        # Optional: Visualize cluster assignment as bar
        matplotlib.pyplot.figure(figsize=(6, 1.5))
        matplotlib.pyplot.title(f"Cluster assignments (k={k})")
        matplotlib.pyplot.plot(labels, 'o-', markersize=8, label='cluster id')
        matplotlib.pyplot.yticks(np.arange(k))
        matplotlib.pyplot.xlabel("Node index")
        matplotlib.pyplot.tight_layout()
        matplotlib.pyplot.savefig(f"plots/demo1_k{k}.png")
        print(f"Saved: plots/demo1_k{k}.png")


if __name__ == '__main__':
    run_demo1()
    # Uncomment to compare with sklearn.cluster.spectral_clustering -- With normalized Laplacian btw!
    print("")
    for k in [2, 3, 4]:
       compare_with_sklearn(k, False)