import cv2
import numpy as np
from yoloface import face_analysis

face_detector = face_analysis()


def equalize_image(data: np.ndarray):
  data_hsv = cv2.cvtColor(data, cv2.COLOR_RGB2HSV)
  data_hsv[:, :, 2] = cv2.equalizeHist(data_hsv[:, :, 2])
  return cv2.cvtColor(data_hsv, cv2.COLOR_HSV2RGB)


def find_face_bbox_yolo(data: np.ndarray):
  _, box, conf = face_detector.face_detection(frame_arr=data, frame_status=True, model='full')
  if len(box) < 1:
    return None, None
  return box, conf


def find_face_bbox(data: np.ndarray):
  classifier_files = [
    'haarcascades/haarcascade_frontalface_default.xml',
    'haarcascades/haarcascade_frontalface_alt.xml',
    'haarcascades/haarcascade_frontalface_alt2.xml',
    'haarcascades/haarcascade_profileface.xml',
    'haarcascades/haarcascade_glasses.xml',
    'lbpcascade_animeface.xml',
  ]
  data_equalized = equalize_image(data)
  data_gray = cv2.cvtColor(data_equalized, cv2.COLOR_RGB2GRAY)
  face_coords, conf = find_face_bbox_yolo(cv2.cvtColor(data_equalized, cv2.COLOR_RGB2BGR))
  if face_coords is not None:
    return face_coords[0]

  for classifier in classifier_files:
    face_cascade = cv2.CascadeClassifier(classifier)
    face_coords = face_cascade.detectMultiScale(data_gray, 1.1, 3)
    if face_coords is not None:
      break
  return max(face_coords, key=lambda v: v[2]*v[3])


def crop_face(data: np.ndarray, bounding_box) -> np.ndarray:
  x, y, w, h = bounding_box
  # Extending the boxes
  factor = 0.4
  x, y = round(x - factor * w), round(y - factor * h)
  w, h = round(w + factor * w * 2), round(h + factor * h * 2)
  y = max(y, 0)
  x = max(x, 0)

  face = data[y:y + h, x:x + w]
  return face