import os
import cv2
from pathlib import Path

# =========================
# CONFIG
# =========================

INPUT_FOLDER = "../Celebrity"     # root folder with subfolders
OUTPUT_FOLDER = "../Celebrity_croped"  # output folder

# extra padding around detected face
PADDING_PERCENT = 0.35

# final output image size
OUTPUT_SIZE = 224

# supported image extensions
IMAGE_EXTENSIONS = [".jpg", ".jpeg", ".png", ".webp"]

# =========================
# LOAD FACE DETECTOR
# =========================

face_cascade = cv2.CascadeClassifier(
    cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
)

# =========================
# HELPERS
# =========================

def make_square_crop(img, x, y, w, h, padding=0.3):
    """
    Create square crop around face with padding.
    """

    img_h, img_w = img.shape[:2]

    # face center
    cx = x + w // 2
    cy = y + h // 2

    # make square size
    side = int(max(w, h) * (1 + padding * 2))

    x1 = max(cx - side // 2, 0)
    y1 = max(cy - side // 2, 0)

    x2 = min(x1 + side, img_w)
    y2 = min(y1 + side, img_h)

    # adjust if crop hits boundaries
    crop_w = x2 - x1
    crop_h = y2 - y1

    side = min(crop_w, crop_h)

    x2 = x1 + side
    y2 = y1 + side

    return img[y1:y2, x1:x2]


def detect_largest_face(img):
    """
    Detect largest face in image.
    """

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    faces = face_cascade.detectMultiScale(
        gray,
        scaleFactor=1.1,
        minNeighbors=5,
        minSize=(50, 50)
    )

    if len(faces) == 0:
        return None

    # choose largest face
    largest = max(faces, key=lambda f: f[2] * f[3])

    return largest


def process_image(input_path, output_path):
    """
    Crop face and save.
    """

    img = cv2.imread(str(input_path))

    if img is None:
        print(f"Failed to read: {input_path}")
        return

    face = detect_largest_face(img)

    if face is None:
        print(f"No face detected: {input_path}")
        return

    x, y, w, h = face

    crop = make_square_crop(
        img,
        x,
        y,
        w,
        h,
        padding=PADDING_PERCENT
    )

    crop = cv2.resize(crop, (OUTPUT_SIZE, OUTPUT_SIZE))

    output_path.parent.mkdir(parents=True, exist_ok=True)

    cv2.imwrite(str(output_path), crop)

    print(f"Saved: {output_path}")


# =========================
# MAIN
# =========================

def main():

    input_root = Path(INPUT_FOLDER)
    output_root = Path(OUTPUT_FOLDER)

    image_files = []

    for ext in IMAGE_EXTENSIONS:
        image_files.extend(input_root.rglob(f"*{ext}"))
        image_files.extend(input_root.rglob(f"*{ext.upper()}"))

    print(f"Found {len(image_files)} images")

    for img_path in image_files:

        relative_path = img_path.relative_to(input_root)
        output_path = output_root / relative_path

        process_image(img_path, output_path)

    print("\nDone.")


if __name__ == "__main__":
    main()