import os
import cv2
import numpy as np
from face_matcher.preprocessing import face_detector, face_aligner
import logging

logging.basicConfig(
    level=logging.ERROR,
    format="{asctime}   {levelname:<8} {message}",
    style="{",
    filename="preprocessor.log",
    filemode="a",
)


class Preprocessor:
    """
    Preprocesses an individual image by detecting, aligning, and cropping (according to a predefined
    margin) the face with the highest score (if any), and preprocesses a dataset splits by going
    over all subject subdirectories and preprocessing all their images.
    """

    def __init__(self, margin=0.2):
        """
        Constructs Preprocessor.

        Attributes
        ----------
        margin: float
            margin to be added to face bounding box height and width (percent decimal of the
            original bounding box size), the margin is doubled during detection and halved
            after face alignment to remove black edges

        Methods
        --------
        preprocessImage(self, img)
            detects, aligns, and crops (according to a predefined margin) the face with the highest
            score in the input image (if any)

        preprocessDirectory(self, source_dir, destination_dir)
            checks subject subdirectories in the source directory in sorted order, finds the
            missing ones from the destination directory, preprocesses the images in those missing
            directories and add them to the destination directory

        testSelf(self)
            tests preprocessor class
        """
        self.margin = margin
        # Initialize face detector
        self.face_detector = face_detector.FaceDetector()
        # Initialize face aligner
        self.face_aligner = face_aligner.FaceAligner()

    def preprocessImage(self, img):
        """
        Detects, aligns, and crops (according to a predefined margin) the face with the highest
        score in the input image (if any).

        Parameters
        ----------
        img: numpy array
            BGR image

        Return
        ------
        img: numpy array or None
            cropped BGR image of the aligned face with the highest score, otherwise it's None if
            no face was detected

        meta: dictionary
            metadata may include the following key/value pairs:
            input_image_size: height and width of the input image
            number_of_faces: number of faces detected in the input image
            face_size: height and width of the detected face image area with the highest
                confidence score (if any)
        """
        # Pad the input image to make it square (if it is not already)
        input_height, input_width, channels = img.shape
        color = (0, 0, 0)
        if input_height > input_width:
            padding_width = round((input_height - input_width) / 2)
            padding_area = np.full(
                (input_height, padding_width, channels), color, dtype=np.uint8
            )
            img = np.concatenate((padding_area, img, padding_area), axis=1)
        elif input_height < input_width:
            padding_height = round((input_width - input_height) / 2)
            padding_area = np.full(
                (padding_height, input_width, channels), color, dtype=np.uint8
            )
            img = np.concatenate((padding_area, img, padding_area), axis=0)

        # # VISUAL DEBUG (squared input image)
        # cv2.imshow("image", img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()

        scores, keypoints, boxes = self.face_detector.detect(img)
        # Check if no face was detected
        if len(scores) == 0:
            meta = {
                "input_image_size": {"height": input_height, "width": input_width},
                "number_of_faces": 0,
            }
            img = None
            return img, meta
        else:
            no_of_faces = len(scores)

            # Only keep detection of the highest score face
            scores, keypoints, boxes = scores[:1], keypoints[:1], boxes[:1]

            # Padded box (bounding box + desired margin)
            box_side_len = boxes[0][2] - boxes[0][0]
            box_side_pixels = box_side_len * img.shape[1]
            box_padding = box_side_len * self.margin
            padded_box = [
                boxes[0][0] - box_padding,
                boxes[0][1] - box_padding,
                boxes[0][2] + box_padding,
                boxes[0][3] + box_padding,
            ]
            top_left_point = [
                round(padded_box[0] * img.shape[1]),
                round(padded_box[1] * img.shape[0]),
            ]
            bottom_right_point = [
                round(padded_box[2] * img.shape[1]),
                round(padded_box[3] * img.shape[0]),
            ]

            # Ensure padded box is still square after projection to pixels (rounding drift)
            if (
                bottom_right_point[0] - top_left_point[0]
                > bottom_right_point[1] - top_left_point[1]
            ):
                bottom_right_point[0] -= 1
            if (
                bottom_right_point[0] - top_left_point[0]
                < bottom_right_point[1] - top_left_point[1]
            ):
                bottom_right_point[1] -= 1
            padded_box_points = [top_left_point, bottom_right_point]

            # # VISUAL DEBUG (bounding box)
            # img_plot = img.copy()
            # img_plot = cv2.rectangle(img_plot, top_left_point, bottom_right_point, (0, 255, 0), 4)
            # cv2.imshow("detections", img_plot)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            # Make sure no negative index and crop face image
            img = img[
                max(padded_box_points[0][1], 0) : padded_box_points[1][1],
                max(padded_box_points[0][0], 0) : padded_box_points[1][0],
            ]

            # Pad face image to make it square in case face is out of frame
            height, width, channels = img.shape
            color = (0, 0, 0)
            if height > width:
                padding_width = round((height - width) / 2)
                padding_area = np.full(
                    (height, padding_width, channels), color, dtype=np.uint8
                )
                img = np.concatenate((padding_area, img, padding_area), axis=1)
            elif height < width:
                padding_height = round((width - height) / 2)
                padding_area = np.full(
                    (padding_height, width, channels), color, dtype=np.uint8
                )
                img = np.concatenate((padding_area, img, padding_area), axis=0)

            # # VISUAL DEBUG (squared face image)
            # cv2.imshow("face", img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            # Align face
            img = self.face_aligner.align(img, keypoints[0])

            # # VISUAL DEBUG (aligned face image)
            # cv2.imshow("aligned face", img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            # Total added pixels divided by 2 to remove half and by another 2 for each side
            crop_margin = round(img.shape[0] * (1 - 1 / (1 + 2 * self.margin)) / 4)
            # Crop black edges
            img = img[crop_margin:-crop_margin, crop_margin:-crop_margin]

            # # VISUAL DEBUG (cropped face image)
            # cv2.imshow("cropped face", img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            face_img_height, face_img_width, _ = img.shape
            meta = {
                "input_image_size": {"height": input_height, "width": input_width},
                "number_of_faces": no_of_faces,
                "face_size": {"height": face_img_height, "width": face_img_width},
                "face_score": float(scores),
            }
            return img, meta

    def preprocessDirectory(self, source_dir, destination_dir):
        """
        Checks subject subdirectories in the source directory in sorted order, finds the missing
        ones from the destination directory, preprocesses the images in those missing directories
        and add them to the destination directory.

        Parameters
        ----------
        source_dir: string
            path of the directory that contains subject subdirectories (e.g. train split)

        destination_dir: string
            path of the directory where the preprocessed subject subdirectories to be saved
        """
        source_sub_dirs = sorted(os.listdir(source_dir))

        for sub_dir in source_sub_dirs:
            destination_sub_dir_path = os.path.join(destination_dir, sub_dir)
            # Create subject subdirectory if it doesn't exist in destination directory
            if not os.path.exists(destination_sub_dir_path):
                os.makedirs(destination_sub_dir_path, exist_ok=True)
                # Preprocess all images in source subject subdirectory
                source_sub_dir_path = os.path.join(source_dir, sub_dir)
                images = sorted(os.listdir(source_sub_dir_path))
                for image in images:
                    try:
                        source_image_path = os.path.join(source_sub_dir_path, image)
                        img = cv2.imread(source_image_path)
                        cropped_img, meta = self.preprocessImage(img)
                        # print(meta)
                        # Save cropped face images to destination subject subdirectory
                        if cropped_img is not None:
                            # Append confidence score and cropped face image size (square) to image name
                            image_name = ".".join(image.split(".")[:-1])
                            image_extension = image.split(".")[-1]
                            new_image_name = (
                                f"{image_name}_"
                                f"{str(round(meta['face_score'], 2))}_"
                                f"{str(meta['face_size']['height'])}."
                                f"{image_extension}"
                            )
                            destination_image_path = os.path.join(
                                destination_sub_dir_path, new_image_name
                            )
                            # print(destination_image_path)
                            cv2.imwrite(destination_image_path, cropped_img)
                        else:
                            destination_image_path = os.path.join(
                                destination_sub_dir_path, image
                            )
                            cv2.imwrite(destination_image_path, img)
                    except Exception as e:
                        logging.error(
                            f"'preprocessing image {image} of subject {sub_dir} failed'",
                            exc_info=True,
                        )

    def testSelf(self):
        """Tests Preprocessor class."""
        # preprocessImage test
        image_path = r"test_data/split/s000/000_no_face.jpg"
        img = cv2.imread(image_path)
        img, meta = self.preprocessImage(img)
        assert img is None
        assert len(meta) == 2
        assert meta["input_image_size"] == {"height": 732, "width": 1005}, meta[
            "input_image_size"
        ]
        assert meta["number_of_faces"] == 0
        image_path = r"test_data/split/s001/002_faces.png"
        img = cv2.imread(image_path)
        img, meta = self.preprocessImage(img)
        assert img.shape == (245, 245, 3), f"cropped image shape: {img.shape}"
        img_hash = hash(tuple(img.flatten()))
        assert (
            img_hash == -4519478635357461561
        ), f"cropped image hash doesn't match: {img_hash}"
        assert len(meta) == 4
        assert meta["input_image_size"] == {"height": 857, "width": 691}, meta[
            "input_image_size"
        ]
        assert meta["number_of_faces"] == 3
        assert meta["face_size"] == {"height": 245, "width": 245}, meta["face_size"]
        assert abs(meta["face_score"] - 0.8938881754875183) < 1e-4, meta["face_score"]

        # preprocessDirectory test
        source_dir = r"test_data/split"
        destination_dir = r"test_data/preprocessed_split"
        self.preprocessDirectory(source_dir, destination_dir)
        assert len(os.listdir(destination_dir)) == 3, {os.listdir(destination_dir)}
        assert len(os.listdir(os.path.join(destination_dir, "s000"))) == 0
        assert len(os.listdir(os.path.join(destination_dir, "s001"))) == 2
        # preprocessDirectory test cleanup
        os.rmdir(os.path.join(destination_dir, "s000"))
        sub_dir = os.path.join(destination_dir, "s001")
        os.remove(
            os.path.join(
                sub_dir, "001_48s4212n-n012-43b2-hs23-847h66nf7347_0.95_1494.jpg"
            )
        )
        os.remove(os.path.join(sub_dir, "002_faces_0.89_245.png"))
        os.rmdir(sub_dir)

        return True
