#
# Demo 3c: Recursive normalized cuts (full version)
#
# author: Christos Choutouridis <cchoutou@ece.auth.gr>
# date:   06/07/2025
#

try:
    from scipy.io import loadmat
    import matplotlib.pyplot as plt
    import numpy as np
    #Project requirements
    from image_to_graph import image_to_graph
    from normalized_cuts import n_cuts_recursive
except ImportError as e:
    print("Missing package:", e)
    print("Run: pip install -r requirements.txt")
    exit(1)

def plot_recursive_clusters(image, labels, title, fname):
    M, N, _ = image.shape
    segmented = labels.reshape(M, N)
    plt.imshow(segmented, cmap='tab20')  # tab20 supports up to 20 unique colors
    plt.title(title)
    plt.axis('off')
    plt.tight_layout()
    plt.savefig(fname)
    print(f"Saved: {fname}")
    plt.close()


def run_demo3c(T1: float, T2: float):
    data = loadmat("dip_hw_3.mat")

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

        affinity = image_to_graph(img)
        labels = n_cuts_recursive(affinity, T1=T1, T2=T2)

        num_clusters = len(np.unique(labels))
        print(f"  Clusters found: {num_clusters}")
        print(f"  Labels: {np.unique(labels)}")

        plot_recursive_clusters(
            img,
            labels,
            title=f"{name} - recursive n_cuts (T1={T1}, T2={T2})",
            fname=f"plots/demo3c_{name}_recursive_T1-{T1}_T2-{T2}.png"
        )


if __name__ == '__main__':
    run_demo3c(5, 0.2)
    run_demo3c(5, 0.95)
    run_demo3c(5, 0.975)