import torch
import os
import random
from PIL import Image
from torchvision.utils import make_grid, save_image
from torchvision import transforms
from tqdm import tqdm

# ==========================================
# 1. CẤU HÌNH
# ==========================================
real_folder = "cifar10_train_ref"
awwl_folder = "generated_images/rescue_detail_a0.2_p1.0_db1"

GRID_SIZE = (8, 8) # Lưới 8x8 = 64 ảnh
PADDING = 2        # Khoảng cách giữa các ảnh
PAD_VALUE = 1.0    # Màu nền padding (1.0 = Trắng, 0.0 = Đen)

# ==========================================
# 2. HÀM TẠO GRID
# ==========================================
def create_grid_image(folder_path, output_name):
    print(f"Creating grid for {folder_path}...")
    
    if not os.path.exists(folder_path):
        print(f"❌ Error: Folder not found {folder_path}")
        return

    # Lấy danh sách ảnh
    files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith('.png')]
    
    # Random sample đủ số lượng
    num_images = GRID_SIZE[0] * GRID_SIZE[1]
    if len(files) < num_images:
        print(f"⚠️ Warning: Not enough images. Need {num_images}, found {len(files)}")
        selected_files = files
    else:
        selected_files = random.sample(files, num_images)
    
    # Load ảnh và chuyển sang Tensor
    to_tensor = transforms.ToTensor()
    tensors = []
    
    for f in selected_files:
        img = Image.open(f).convert("RGB")
        tensors.append(to_tensor(img))
    
    # Xếp thành batch [B, C, H, W]
    batch = torch.stack(tensors)
    
    # Tạo lưới
    grid = make_grid(batch, nrow=GRID_SIZE[0], padding=PADDING, pad_value=PAD_VALUE)
    
    # Lưu ảnh
    save_image(grid, output_name)
    print(f"✅ Saved {output_name}")

# ==========================================
# 3. CHẠY
# ==========================================
if __name__ == "__main__":
    # Tạo lưới Real Data
    create_grid_image(real_folder, "grid_real_samples.png")
    
    # Tạo lưới AWWL
    create_grid_image(awwl_folder, "grid_awwl_generated.png")