""" 
Before running this script, please put the RainDS/RainDS_real folder in the same directory as this script,
then rename the folder to 'dataset'.
This script will remove any duplicate scene names (png files) that exist in both the training and test sets.
"""

import os

base_dir = "dataset"

train_gt_dir = os.path.join(base_dir, "train_set", "gt")
test_dirs = {
    "gt": os.path.join(base_dir, "test_set", "gt"),
    "raindrop": os.path.join(base_dir, "test_set", "raindrop"),
    "rainstreak": os.path.join(base_dir, "test_set", "rainstreak"),
    "rainstreak_raindrop": os.path.join(base_dir, "test_set", "rainstreak_raindrop"),
}

# Get all scene names (png files) in train GT folder
train_files = set(f for f in os.listdir(train_gt_dir) if f.endswith(".png"))
print(train_files)

# Get all scene names (png files) in test GT folder
test_gt_dir = test_dirs["gt"]
test_files = set(f for f in os.listdir(test_gt_dir) if f.endswith(".png"))

# Find intersection
duplicate_files = train_files.intersection(test_files)

for filename in sorted(duplicate_files, key=lambda x: int(x.split(".")[0])):
    for key, path in test_dirs.items():
        file_path = os.path.join(path, filename)
        if os.path.exists(file_path):
            print(f"removing {filename} from dataset/test_set/{key}")
            os.remove(file_path)