import json
import numpy as np
import os

def symlink(json_path, save_path, 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]
    os.makedirs(os.path.join(save_path, mode), exist_ok=True)
    for img_path in worst_images:
        if not os.path.exists(os.path.join(save_path, mode, os.path.basename(img_path))):
            os.symlink(img_path, os.path.join(save_path, mode, os.path.basename(img_path)))
    return len(worst_images)

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"]:
        num_symlinks = symlink(f'{ROOT_FOLDER}/{subfolder}/aesthetic.json', f'{ROOT_FOLDER}/{subfolder}', mode=mode)
        print(f"Created {num_symlinks} symlinks for {subfolder} in {mode} mode.")
