import os
from abc import ABC, abstractmethod
from typing import List

import cv2
import numpy as np
from retinaface import RetinaFace
from retinaface.model import retinaface_model

from .box_utils import convert_to_square


class FaceDetector(ABC):
    def __init__(self, target_size):
        self.target_size = target_size
    @abstractmethod
    def detect_crops(self, img, *args, **kwargs) -> List[np.ndarray]:
        """
        Img is a numpy ndarray in range [0..255], uint8 dtype, RGB type
        Returns ndarray with [x1, y1, x2, y2] in row
        """
        pass

    @abstractmethod
    def postprocess_crops(self, crops, *args, **kwargs) -> List[np.ndarray]:
        return crops

    def sort_faces(self, crops):
        sorted_faces = sorted(crops, key=lambda x: -(x[2] - x[0]) * (x[3] - x[1]))
        sorted_faces = np.stack(sorted_faces, axis=0)
        return sorted_faces

    def fix_range_crops(self, img, crops):
        H, W, _ = img.shape
        final_crops = []
        for crop in crops:
            x1, y1, x2, y2 = crop
            x1 = max(min(round(x1), W), 0)
            y1 = max(min(round(y1), H), 0)
            x2 = max(min(round(x2), W), 0)
            y2 = max(min(round(y2), H), 0)
            new_crop = [x1, y1, x2, y2]
            final_crops.append(new_crop)
        final_crops = np.array(final_crops, dtype=np.int)
        return final_crops

    def crop_faces(self, img, crops) -> List[np.ndarray]:
        cropped_faces = []
        for crop in crops:
            x1, y1, x2, y2 = crop
            face_crop = img[y1:y2, x1:x2, :]
            cropped_faces.append(face_crop)
        return cropped_faces

    def unify_and_merge(self, cropped_images):
        return cropped_images

    def __call__(self, img):
        return self.detect_faces(img)

    def detect_faces(self, img):
        crops = self.detect_crops(img)
        if crops is None or len(crops) == 0:
            return [], []
        crops = self.sort_faces(crops)
        updated_crops = self.postprocess_crops(crops)
        updated_crops = self.fix_range_crops(img, updated_crops)
        cropped_faces = self.crop_faces(img, updated_crops)
        unified_faces = self.unify_and_merge(cropped_faces)
        return unified_faces, updated_crops


class StatRetinaFaceDetector(FaceDetector):
    def __init__(self, target_size=None):
        super().__init__(target_size)
        self.model = retinaface_model.build_model()
        #self.relative_offsets = [0.3258, 0.5225, 0.3258, 0.1290]
        self.relative_offsets = [0.3619, 0.5830, 0.3619, 0.1909]

    def postprocess_crops(self, crops, *args, **kwargs) -> np.ndarray:
        final_crops = []
        x1_offset, y1_offset, x2_offset, y2_offset = self.relative_offsets
        for crop in crops:
            x1, y1, x2, y2 = crop
            w, h = x2 - x1, y2 - y1
            x1 -= w * x1_offset
            y1 -= h * y1_offset
            x2 += w * x2_offset
            y2 += h * y2_offset
            crop = np.array([x1, y1, x2, y2], dtype=crop.dtype)
            crop = convert_to_square(crop)
            final_crops.append(crop)
        final_crops = np.stack(final_crops, axis=0)
        return final_crops

    def detect_crops(self, img, *args, **kwargs):
        faces = RetinaFace.detect_faces(img, model=self.model)
        crops = []
        if isinstance(faces, tuple):
            faces = {}
        for name, face in faces.items():
            x1, y1, x2, y2 = face['facial_area']
            crop = np.array([x1, y1, x2, y2])
            crops.append(crop)
        if len(crops) > 0:
            crops = np.stack(crops, axis=0)
        return crops

    def unify_and_merge(self, cropped_images):
        if self.target_size is None:
            return cropped_images
        else:
            resized_images = []
            for cropped_image in cropped_images:
                resized_image = cv2.resize(cropped_image, (self.target_size, self.target_size),
                                           interpolation=cv2.INTER_LINEAR)
                resized_images.append(resized_image)

            resized_images = np.stack(resized_images, axis=0)
            return resized_images