import torch
import torch.nn.functional as F
import os
import subprocess

CV2_AVAILABLE = True
try:
    import cv2
except:
    print("OpenCV is not installed so face cropping is not available.")
    CV2_AVAILABLE = False

CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
DETECTOR_FILE = "lbpcascade_animeface.xml"

if not os.path.exists(os.path.join(CURRENT_DIR, DETECTOR_FILE)):
    print("Downloading anime face detector...")
    try:
        subprocess.run(["wget", "https://raw.githubusercontent.com/nagadomi/lbpcascade_animeface/master/lbpcascade_animeface.xml", "-P", CURRENT_DIR])
    except:
        print(f"Failed to download lbpcascade_animeface.xml so please download it in {CURRENT_DIR}.")
        CV2_AVAILABLE = False

CROP_MODES = ["padding", "face_crop", "none"] if CV2_AVAILABLE else ["padding", "none"]

def image_to_numpy(image):
    image = image.squeeze(0) * 255
    return image.numpy().astype("uint8")

def numpy_to_image(image):
    image = torch.tensor(image).float() / 255
    return image.unsqueeze(0)

def pad_to_square(tensor):
    tensor = tensor.squeeze(0).permute(2, 0, 1)
    _, h, w = tensor.shape

    target_length = max(h, w)

    pad_l = (target_length - w) // 2
    pad_r = (target_length - w) - pad_l
    
    pad_t = (target_length - h) // 2
    pad_b = (target_length - h) - pad_t

    padded_tensor = F.pad(tensor, (pad_l, pad_r, pad_t, pad_b), mode="constant", value=0)

    return padded_tensor.permute(1, 2, 0).unsqueeze(0)

def face_crop(image):
    image = image_to_numpy(image)
    face_cascade = cv2.CascadeClassifier(os.path.join(CURRENT_DIR, DETECTOR_FILE))
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)

    w, h = image.shape[1], image.shape[0]

    target_length = min(w, h)
    fx, fy, fw, fh = (0, 0, w, h) if len(faces) == 0 else faces[0]

    dx = target_length - fw // 2
    dy = target_length - fh // 2

    target_x = 0 if w < h else max(0, fx - dx)
    target_y = 0 if w > h else max(0, fy - dy)
    
    image = image[target_y:target_y+target_length, target_x:target_x+target_length]
    image = numpy_to_image(image)

    return image

class ImageCrop:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE", ),
                "mode": (CROP_MODES, ), 
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "preprocess"
    CATEGORY = "image/preprocessors"

    def preprocess(self, image, mode):
        if mode == "padding":
            image = pad_to_square(image) 
        elif mode == "face_crop":
            image = face_crop(image)
        
        return (image,)