"""
Visualize sample images and masks from the CMP Facade dataset.
"""
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from PIL import Image
import random


def visualize_samples(data_dir="data/cmp_facade", split="train", num_samples=5, save_path="outputs/dataset_visualization.png"):
    os.makedirs(os.path.dirname(save_path), exist_ok=True)

    split_dir = f"{data_dir}/{split}"
    img_dir = f"{split_dir}/images"
    mask_dir = f"{split_dir}/masks_binary"

    img_files = sorted([f for f in os.listdir(img_dir) if f.endswith('.jpg')])[:20]
    selected = random.sample(img_files, min(num_samples, len(img_files)))

    fig, axes = plt.subplots(num_samples, 3, figsize=(12, 4 * num_samples))
    if num_samples == 1:
        axes = axes.reshape(1, -1)

    for i, img_file in enumerate(selected):
        img = Image.open(f"{img_dir}/{img_file}").convert('RGB')
        mask = Image.open(f"{mask_dir}/{img_file.replace('.jpg', '.png')}").convert('L')
        mask_arr = np.array(mask)

        img_arr = np.array(img)
        overlay = img_arr.copy()
        overlay[mask_arr > 0] = [255, 0, 0]
        blended = (img_arr * 0.6 + overlay * 0.4).astype(np.uint8)

        axes[i, 0].imshow(img)
        axes[i, 0].set_title(f"Image {img_file}")
        axes[i, 0].axis('off')

        axes[i, 1].imshow(mask, cmap='gray')
        axes[i, 1].set_title("Facade Mask (class 1)")
        axes[i, 1].axis('off')

        axes[i, 2].imshow(blended)
        axes[i, 2].set_title("Overlay")
        axes[i, 2].axis('off')

    plt.tight_layout()
    plt.savefig(save_path, dpi=150)
    plt.close()
    print(f"Visualization saved to {save_path}")


if __name__ == "__main__":
    visualize_samples("data/cmp_facade", "train", num_samples=5, save_path="outputs/dataset_visualization.png")