import json
import numpy as np
import random
import os

from PIL import Image, ImageOps
import math
def visualize(json_path, save_path, size=512, mode="worst"):
    with open(json_path, 'r') as f:
        data = json.load(f)


    scores = np.array(list(data.values()))
    if mode == "worst":
        threshold = np.percentile(scores, 10)  # 10th percentile
        worst_images = [img for img, score in data.items() if score <= threshold]
    elif mode == "best":
        threshold = np.percentile(scores, 90)
        worst_images = [img for img, score in data.items() if score >= threshold]

    sampled_images = random.sample(worst_images, 9)

    def pad_image(img, desired_size):
        old_size = img.size
        ratio = float(desired_size) / max(old_size)
        new_size = tuple([int(x * ratio) for x in old_size])
        img = img.resize(new_size, Image.Resampling.BICUBIC)
        delta_w = desired_size - new_size[0]
        delta_h = desired_size - new_size[1]
        padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
        new_img = ImageOps.expand(img, padding)
        return new_img
    images = []
    desired_size = size  # You can adjust this size as needed

    for img_path in sampled_images:
        img = Image.open(img_path)
        img = pad_image(img, desired_size)
        images.append(img)
    grid_cols = 3
    grid_rows = math.ceil(len(images) / grid_cols)
    grid_width = grid_cols * desired_size
    grid_height = grid_rows * desired_size
    grid_image = Image.new('RGB', (grid_width, grid_height))
    for idx, img in enumerate(images):
        x = (idx % grid_cols) * desired_size
        y = (idx // grid_cols) * desired_size
        grid_image.paste(img, (x, y))
    grid_image.save(save_path)
    # or display it
    #grid_image.show()

ROOT_FOLDER = "/ROOT"
for subfolder in os.listdir(ROOT_FOLDER):
    if not os.path.isdir(f'{ROOT_FOLDER}/{subfolder}'):
        continue
    for mode in ["worst", "best"]:
        visualize(f'{ROOT_FOLDER}/{subfolder}/aesthetic.json', f'{ROOT_FOLDER}/{subfolder}_aesthetic_{mode}.jpg', mode=mode)
