import random
import matplotlib.pyplot as plt

# -----------------------------
# Pick random samples
# -----------------------------
num_samples = 19
# indices = random.sample(range(len(mnist_original)), num_samples)

# def visualize_sample(dataset, num_samples):
#     indices = random.sample(range(len(dataset)), num_samples)

#     # -----------------------------
#     # Plot 1 x num_samples row
#     # -----------------------------
#     fig, axes = plt.subplots(1, num_samples, figsize=(num_samples * 2, 2.5), squeeze=False)

#     for col, idx in enumerate(indices):
#         img, label = dataset[idx]

#         axes[0, col].imshow(img.squeeze(), cmap="gray")
#         axes[0, col].set_title(str(label))
#         axes[0, col].axis("off")

#     plt.tight_layout()
#     plt.show()

def visualize_sample(mnist_original, mnist_processed, num_samples):
    indices = random.sample(range(len(mnist_original)), num_samples)

    # -----------------------------
    # Plot 2 x 8 grid
    # -----------------------------
    fig, axes = plt.subplots(2, num_samples, figsize=(num_samples * 2, 4))

    for col, idx in enumerate(indices):
        orig_img, label = mnist_original[idx]
        proc_img, _ = mnist_processed[idx]

        # Top row: original
        axes[0, col].imshow(orig_img.squeeze(), cmap="gray")
        axes[0, col].set_title(str(label))
        axes[0, col].axis("off")

        # Bottom row: processed
        axes[1, col].imshow(proc_img.squeeze(), cmap="gray")
        axes[1, col].axis("off")

    axes[0, 0].set_ylabel("Original\n28×28", fontsize=10)
    axes[1, 0].set_ylabel("Processed\n1-bit 5×5", fontsize=10)

    plt.tight_layout()
    plt.show()