from pathlib import Path
from PIL import Image
import shutil
import statistics

# --- Settings ---
INPUT_DIR = Path("input_images")        # where all images are
OUTPUT_DIR = Path("input_square")       # where selected square-ish images will be copied
ASPECT_MAX = 1.20                       # max ratio between longer/shorter side (1.00 = perfect square)
SIZE_TOLERANCE = 0.35                   # +-35% around median "short side" (set None to disable)

exts = {".png", ".jpg", ".jpeg", ".webp"}

OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

candidates = []
short_sides = []

# --- Scan images and compute metrics ---
files = sorted([p for p in INPUT_DIR.iterdir() if p.is_file() and p.suffix.lower() in exts])
print(f"Found {len(files)} files in {INPUT_DIR.resolve()}")

for p in files:
    with Image.open(p) as im:
        w, h = im.size
    long_side = max(w, h)
    short_side = min(w, h)
    aspect = long_side / short_side

    # Keep only square-ish by aspect ratio
    if aspect <= ASPECT_MAX:
        candidates.append((p, w, h, short_side, aspect))
        short_sides.append(short_side)

print(f"Square-ish (aspect <= {ASPECT_MAX}): {len(candidates)}")

# --- Optional: filter also by similar size (short side around median) ---
if SIZE_TOLERANCE is not None and len(short_sides) > 0:
    med = statistics.median(short_sides)
    lo = med * (1 - SIZE_TOLERANCE)
    hi = med * (1 + SIZE_TOLERANCE)

    filtered = [c for c in candidates if lo <= c[3] <= hi]
    print(f"After size filter (median short side={med:.0f}, range {lo:.0f}-{hi:.0f}): {len(filtered)}")
else:
    filtered = candidates

# --- Copy selected to output folder ---
copied = 0
for (p, w, h, short_side, aspect) in filtered:
    dst = OUTPUT_DIR / p.name
    shutil.copy2(p, dst)
    copied += 1

print(f"Copied {copied} images to {OUTPUT_DIR.resolve()}")