#!/usr/bin/env python3
"""
GPU YOLO photo filter: keep people/scenic/screenshots, move everything else.
Run with: python3 -u yolo_no_people_filter.py   (-u for unbuffered output)

Customize SRC, DST, and class_sets below.

When running via Hermes background processes, ALWAYS use python3 -u and flush=True
on all print() calls. Without this, progress output stays buffered and won't appear
in the process viewer until the script exits.
"""

import os, shutil
from pathlib import Path
from collections import Counter
import torch
from ultralytics import YOLO

# === CONFIGURATION ===
SRC = "/mnt/seagate8tb/Photos"       # Source directory (recursive)
DST = "/mnt/seagate8tb/NO PEOPLE"    # Destination for filtered photos
BATCH_SIZE = 2                        # 2 for 4GB GPU, 4 for 8GB+
NUM_WORKERS = 4
DEVICE = "cuda:0"
CONFIDENCE = 0.3
MODEL = "yolov8s.pt"                 # s=11M params (4GB safe), n=3M, m=26M
HALF = True                          # FP16 for speed on CUDA

# === CLASSIFICATION SETS (COCO indices) ===
PEOPLE_CLASSES = {0}

SCREENSHOT_CLASSES = {
    62, 63, 64, 65, 66, 67,  # tv, laptop, mouse, remote, keyboard, cell phone
    73, 76,                  # book, scissors
}

SCENIC_CLASSES = {
    1, 2, 3, 4, 6, 7, 8,    # bicycle, car, motorcycle, airplane, bus, train, truck
    14, 15, 16, 17, 18, 19,   # bird, cat, dog, horse, sheep, cow
    20, 21, 22, 23,           # elephant, bear, zebra, giraffe
    25, 26, 27, 28,           # umbrella, handbag, tie, suitcase
    29, 30, 31, 32, 33, 34,   # frisbee, skis, snowboard, sports ball, kite, baseball bat
    35, 36, 37, 38,           # baseball glove, skateboard, surfboard, tennis racket
    39, 40, 41,               # bottle, wine glass, cup
    42, 43, 44, 45, 46,       # fork, knife, spoon, bowl, banana
    47, 48, 49, 50, 51,       # apple, sandwich, orange, broccoli, carrot
    52, 53, 54, 55,           # hot dog, pizza, donut, cake
    56, 57, 58, 59, 60,       # chair, couch, potted plant, bed, dining table
    72, 74, 75, 77, 78, 79,   # clock, vase, teddy bear, hair drier, toothbrush
    80, 81,                   # hair drier, toothbrush (duplicates from COCO)
}


def is_scenic(detections):
    return any(cid in SCENIC_CLASSES for cid in detections)


def is_screenshot(detections):
    return any(cid in SCREENSHOT_CLASSES for cid in detections)


def main():
    print(f"Loading {MODEL} on {DEVICE}...", flush=True)
    model = YOLO(MODEL)
    model.to(DEVICE)

    # Collect all image files
    extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp',
                  '.tiff', '.tif', '.gif', '.heic', '.heif'}
    files = []
    for root, dirs, fnames in os.walk(SRC):
        for f in fnames:
            if Path(f).suffix.lower() in extensions:
                files.append(os.path.join(root, f))

    total = len(files)
    print(f"Found {total} images to scan", flush=True)

    if total == 0:
        print("No images found!")
        return

    stats = Counter()
    moved = 0
    kept_people = 0
    kept_scenic = 0
    kept_screenshot = 0
    errors = 0

    for i in range(0, total, BATCH_SIZE):
        batch = files[i:i + BATCH_SIZE]
        batch_idx = i // BATCH_SIZE + 1
        total_batches = (total + BATCH_SIZE - 1) // BATCH_SIZE

        try:
            results = model(batch, device=DEVICE, verbose=False,
                          conf=CONFIDENCE, imgsz=640, half=HALF, stream=False)
        except Exception as e:
            # Batch failed — likely corrupt images in the mix.
            # Move them all directly to NO PEOPLE (no GPU retry — too slow).
            for filepath in batch:
                try:
                    rel = os.path.relpath(filepath, SRC)
                    dst = os.path.join(DST, rel)
                    os.makedirs(os.path.dirname(dst), exist_ok=True)
                    shutil.move(filepath, dst)
                    moved += 1
                    stats['no_people'] += 1
                except Exception as e2:
                    errors += 1
            continue

        for result, filepath in zip(results, batch):
            boxes = result.boxes
            if boxes is None or len(boxes) == 0:
                # No detections → move
                rel = os.path.relpath(filepath, SRC)
                dst = os.path.join(DST, rel)
                os.makedirs(os.path.dirname(dst), exist_ok=True)
                shutil.move(filepath, dst)
                moved += 1
                stats['empty'] += 1
                continue

            class_ids = boxes.cls.cpu().numpy().astype(int)
            has_people = any(cid in PEOPLE_CLASSES for cid in class_ids)
            scenic = is_scenic(class_ids)
            screenshot = is_screenshot(class_ids)

            if has_people:
                kept_people += 1
                stats['people'] += 1
            elif screenshot:
                kept_screenshot += 1
                stats['screenshot'] += 1
            elif scenic:
                kept_scenic += 1
                stats['scenic'] += 1
            else:
                rel = os.path.relpath(filepath, SRC)
                dst = os.path.join(DST, rel)
                os.makedirs(os.path.dirname(dst), exist_ok=True)
                try:
                    shutil.move(filepath, dst)
                    moved += 1
                    stats['no_people'] += 1
                except Exception as e:
                    print(f"  Move error: {filepath}: {e}", flush=True)
                    errors += 1

        if batch_idx % 50 == 0 or batch_idx == total_batches:
            elapsed = min(batch_idx * BATCH_SIZE, total)
            pct = elapsed / total * 100
            print(f"  [{elapsed}/{total}] {pct:.1f}% | moved={moved} | "
                  f"people={kept_people} | scenic={kept_scenic} | "
                  f"ss={kept_screenshot} | err={errors}", flush=True)

    print(f"\n{'='*60}")
    print(f"DONE: {total} images processed")
    print(f"  👤 People (kept):      {kept_people}")
    print(f"  🏔️  Scenic (kept):      {kept_scenic}")
    print(f"  📱 Screenshot (kept):   {kept_screenshot}")
    print(f"  📦 Moved to NO PEOPLE:  {moved}")
    print(f"  ❌ Errors:              {errors}")
    print(f"  💾 Empty (no detections): {stats.get('empty', 0)}")


if __name__ == "__main__":
    main()
