"""
Prepare images for COLMAP SfM + 3DGS.
- Convert RGBA -> RGB on white background
- Resize so max dimension <= 1600 (COLMAP works better with reasonable res)
- Organize into per-part input folders: data/parts/<part>/input/*.png
"""
import os
from PIL import Image
import numpy as np

SRC = "data/raw"
DST = "data/parts"
MAX_DIM = 1600

os.makedirs(DST, exist_ok=True)

for part in sorted(os.listdir(SRC)):
    src_dir = os.path.join(SRC, part)
    if not os.path.isdir(src_dir):
        continue
    dst_dir = os.path.join(DST, part, "input")
    os.makedirs(dst_dir, exist_ok=True)
    for fname in sorted(os.listdir(src_dir)):
        if not fname.lower().endswith(".png"):
            continue
        im = Image.open(os.path.join(src_dir, fname)).convert("RGBA")
        arr = np.array(im)
        # composite over white background
        alpha = arr[:, :, 3:4].astype(np.float32) / 255.0
        rgb = arr[:, :, :3].astype(np.float32)
        bg = np.array([255, 255, 255], dtype=np.float32)
        composited = rgb * alpha + bg * (1 - alpha)
        out = Image.fromarray(composited.astype(np.uint8), "RGB")
        # resize
        w, h = out.size
        scale = MAX_DIM / max(w, h)
        if scale < 1.0:
            nw, nh = int(w * scale), int(h * scale)
            out = out.resize((nw, nh), Image.LANCZOS)
        # rename to simple sequential
        idx = fname.split("_")[-1].split(".")[0]
        out_name = f"{part}_{idx}.png"
        out.save(os.path.join(dst_dir, out_name))
        print(f"{part}/{fname} -> {out_name} {out.size}")

print("Done preparing images.")