from gpiozero import Button, LED
from time import sleep, localtime, strftime
from picamera2 import Picamera2
from PIL import Image
import os, subprocess, signal
from signal import pause
from rembg import remove
import cv2
import numpy as np

button_norm = Button(22, pull_up = True)
button_pi = Button(23, pull_up = True)
red = LED(17)
white = LED(27)
home_dir = os.environ['HOME']

cam = Picamera2()
config = cam.create_still_configuration()
cam.configure(config)

def ledcam():
    # Countdown LED
    for i in range(3):
        red.on()
        sleep(0.5)
        red.off()
        sleep(0.5)
    white.on()
    sleep(1)

    # Filename setup
    t = localtime()
    file_name = strftime('%H-%M-%S', t)
    raw_path = f"{home_dir}/Photos/{file_name}.jpg"
    mask_path = f"{home_dir}/norm_mask.jpg"  
    final_path = f"{home_dir}/Photos/{file_name}_final.jpg"

    # Capture photo
    cam.start()
    cam.capture_file(raw_path)
    cam.stop()
    print("Successfully taken photo")
    white.off()

    # MASKING LOGIC
    photo = Image.open(raw_path).convert("RGBA")
    mask = Image.open(mask_path).convert("RGBA")

    photo_width = mask.width
    aspect_ratio = photo.height / photo.width
    photo_height = int(photo_width * aspect_ratio)
    photo = photo.resize((photo_width, photo_height))

    # Create canvas and paste photo onto mask
    mask.paste(photo, (0, 0), photo) 
    
    # Dither
    final_img = mask.convert('L')
    dithered_img = final_img.convert('1', dither=Image.Dither.FLOYDSTEINBERG)
    dithered_img.save(final_path)

    try:
        subprocess.run(["lp", "-d", "SN_DB402", "-o", "fit-to-page", final_path])
        print("Successfully printed masked photo")
    except Exception as e:
        print(f"Error: {e}")


def ledcam_pi():
    # Countdown LED
    for i in range(3):
        red.on()
        sleep(0.5)
        red.off()
        sleep(0.5)
    white.on()
    sleep(1)

    # Filename setup
    t = localtime()
    file_name = strftime('%H-%M-%S', t)
    raw_path = f"{home_dir}/Photos/{file_name}.jpg"
    mask_path = f"{home_dir}/pi_mask.jpg"
    final_path = f"{home_dir}/Photos/{file_name}_final.jpg"

    # Capture photo
    cam.start()
    cam.capture_file(raw_path)
    cam.stop()
    print("Successfully taken photo")
    white.off()
    
    # 1. Load image for detection
    img = cv2.imread(raw_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 2. Face detection
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
    faces = face_cascade.detectMultiScale(gray, 1.1, 4)

    canvas = Image.open(mask_path).convert("RGBA")

    if len(faces) == 0:
        print("No face detected, using original image.")
        photo = Image.open(raw_path).convert("RGBA")
        photo_width = canvas.width
        photo = photo.resize((photo_width, int(photo_width * (photo.height / photo.width))))
        canvas.paste(photo, (0, 0), photo)
    else:
        x, y, w, h = faces[0]
        padding = int(w * 0.2)
        face_crop = img[max(0, y-padding):y+h+padding, max(0, x-padding):x+w+padding]
        
        # Convert to PILLOW and remove background
        face_pil = Image.fromarray(cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB))
        face_pil = remove(face_pil) 
        
        # Scale and Center
        new_size = int(canvas.width * 0.4) 
        face_pil = face_pil.resize((new_size, new_size), Image.Resampling.LANCZOS)
        
        paste_x = (canvas.width - face_pil.width) // 2
        paste_y = (canvas.height - face_pil.height) // 4 
        canvas.paste(face_pil, (paste_x, paste_y), face_pil)

    # Dither and Save
    final_img = canvas.convert('L').convert('1', dither=Image.Dither.FLOYDSTEINBERG)
    final_img.save(final_path) 
        
    try:
        subprocess.run(["lp", "-d", "SN_DB402", "-o", "fit-to-page", final_path])
        print("Successfully printed masked photo")
    except Exception as e:
        print(f"Error: {e}")


button_norm.when_pressed = ledcam
button_pi.when_pressed = ledcam_pi
pause()