"""
prepare_examples.py
-------------------
Utility script: extracts 5 sample CT scan images from the IBM dataset
and saves them as PNG files in the examples/ folder for Gradio demos.

Run once before pushing to Hugging Face:
    python prepare_examples.py
    python prepare_examples.py --data_dir "C:/path/to/2d_images"
"""

import argparse
import os
import glob
import numpy as np
from skimage.io import imread
from skimage.transform import resize
from PIL import Image


def prepare_examples(data_dir=".", n=5, output_dir="examples"):
    """
    Load n CT scan .tif files, convert to PNG, save to output_dir.

    Args:
        data_dir   : directory containing 2d_images/**/*.tif
        n          : number of example images to extract
        output_dir : folder to save PNG files into
    """
    os.makedirs(output_dir, exist_ok=True)

    # Search for .tif files
    pattern = os.path.join(data_dir, "2d_images", "**", "*.tif")
    img_list = sorted(glob.glob(pattern, recursive=True))

    if not img_list:
        # Try the IBMDATA desktop path as fallback
        fallback = r"C:\Users\Shind\OneDrive\Desktop\IBMDATA\2d_images"
        pattern  = os.path.join(fallback, "**", "*.tif")
        img_list = sorted(glob.glob(pattern, recursive=True))

    if not img_list:
        print(
            "ERROR: No .tif images found.\n"
            "Run from the directory that contains 2d_images/, or pass --data_dir."
        )
        return

    # Pick n evenly-spaced images
    indices = np.linspace(0, len(img_list) - 1, min(n, len(img_list)), dtype=int)
    chosen  = [img_list[i] for i in indices]

    print(f"Extracting {len(chosen)} example images to '{output_dir}/'...")
    for k, path in enumerate(chosen, start=1):
        img = imread(path).astype(np.float32)

        # Resize to 256x256 (matches model input)
        img_resized = resize(img, (256, 256), preserve_range=True, anti_aliasing=True)

        # Normalise to uint8 for display
        lo, hi = img_resized.min(), img_resized.max()
        if hi > lo:
            img_uint8 = ((img_resized - lo) / (hi - lo) * 255).astype(np.uint8)
        else:
            img_uint8 = np.zeros((256, 256), dtype=np.uint8)

        out_path = os.path.join(output_dir, f"ct_sample_{k:02d}.png")
        Image.fromarray(img_uint8).save(out_path)
        print(f"  [{k}/{len(chosen)}] Saved: {out_path}")

    print("Done. Example images are ready in", output_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare example CT scan PNGs for Gradio")
    parser.add_argument("--data_dir",   type=str, default=".",       help="Directory with 2d_images/")
    parser.add_argument("--n",          type=int, default=5,         help="Number of examples to extract")
    parser.add_argument("--output_dir", type=str, default="examples", help="Output folder for PNGs")
    args = parser.parse_args()

    prepare_examples(data_dir=args.data_dir, n=args.n, output_dir=args.output_dir)