import os
import cv2
import numpy as np
from skimage import data, img_as_ubyte
from skimage.transform import resize

def prepare_data():
    os.makedirs("test_data", exist_ok=True)

    # Load a standard color image (Astronaut)
    print("Loading Ground Truth image...")
    gt = img_as_ubyte(data.astronaut())
    # Save GT
    cv2.imwrite("test_data/ground_truth.jpg", cv2.cvtColor(gt, cv2.COLOR_RGB2BGR))

    resolutions = {
        "128": (128, 128),
        "512": (512, 512),
        "1080p": (1920, 1080)
    }

    for name, size in resolutions.items():
        print(f"Generating {name}...")
        # Resize GT
        # Note: resize expects float, so we convert back to ubyte
        resized_gt = resize(gt, (size[1], size[0]), anti_aliasing=True) # size is (h, w)
        resized_gt = img_as_ubyte(resized_gt)

        # Save Resized GT
        gt_path = f"test_data/{name}_gt.jpg"
        cv2.imwrite(gt_path, cv2.cvtColor(resized_gt, cv2.COLOR_RGB2BGR))

        # Convert to Grayscale
        gray = cv2.cvtColor(resized_gt, cv2.COLOR_RGB2GRAY)

        # Save Grayscale Input
        gray_path = f"test_data/{name}_gray.jpg"
        cv2.imwrite(gray_path, gray)

    print("Data preparation complete.")

if __name__ == "__main__":
    prepare_data()