from PIL import Image
import numpy as np
from face_matcher.preprocessing import preprocessor
from autocapture_model.detector import DocumentDetector
#from face_matcher.preprocessing.detect_document_photo import DocumentDetector
import os


class ImageProcessor:
    def __init__(self, debug_mode=False) -> None:
        os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # Force doc detector to use CPU
        self.document_detector = DocumentDetector()
        self.prp = preprocessor.Preprocessor()
        self.debug_mode = debug_mode

    def __call__(self, image_dicts):
        images = [i["raw_image"] for i in image_dicts]
        for image, image_dict in zip(images, image_dicts):
            if image_dict.get("detect_document_photo", False):
                result = self.detect_photo_box_and_rotation(image)
                if "box" in result:
                    image_dict["crop"] = result["box"]
                if "angle" in result:
                    image_dict["prerotate_counterclockwise"] = result["angle"]

        images_prerotate = [
            (
                i.rotate(meta["prerotate_counterclockwise"], expand=True)
                if "prerotate_counterclockwise" in meta
                else i
            )
            for i, meta in zip(images, image_dicts)
        ]
        images_cropped = [
            i.crop(meta["crop"]) if "crop" in meta else i
            for i, meta in zip(images_prerotate, image_dicts)
        ]
        images_postrotate = [
            (
                i.rotate(meta["postrotate_counterclockwise"], expand=True)
                if "postrotate_counterclockwise" in meta
                else i
            )
            for i, meta in zip(images_cropped, image_dicts)
        ]
        images_np = [np.asarray(i) for i in images_postrotate]

        preprocessor_outputs = [self.prp.preprocessImage(i) for i in images_np]

        if self.debug_mode: # Save cropped images for debugging
            [im.save(f"doc_detector_out_{i}.png") for i, im in enumerate(images_postrotate)]
            [Image.fromarray(prp_out[0]).save(f"face_detector_out_{i}.png") for i, prp_out in enumerate(preprocessor_outputs)]

        for image_dict, image_full, prp_out in zip(
            image_dicts, images_postrotate, preprocessor_outputs
        ):
            image_dict["preprocessor_metadata"] = prp_out[1]
            image_dict["face_detected"] = prp_out[0] is not None
            image_dict["processed_image"] = (
                Image.fromarray(prp_out[0])
                if image_dict["face_detected"]
                else image_full
            )

        return image_dicts

    def box_to_square(self, box):
        square_box = []
        width = box[2] - box[0]
        height = box[3] - box[1]
        x_center = (box[2] + box[0])/2
        y_center = (box[3] + box[1])/2
        if width >= height:
            square_box.append(int(box[0]))
            square_box.append(int((y_center - (width/2))))
            square_box.append(int(box[2]))
            square_box.append(square_box[1] + width)
        else:
            square_box.append(int((x_center - (height/2))))
            square_box.append(int(box[1]))
            square_box.append(square_box[0] + height)
            square_box.append(int(box[3]))
        return square_box

    def detect_photo_box_and_rotation(self, original_image):
        images_dict = {}
        for rotation in [0, 90, 180, 270]:
            images_dict[rotation] = original_image.rotate(rotation, Image.NEAREST, expand=1)

        response = self.document_detector.detect_batch(images_dict)

        best_result = {"score":0, "angle":0, "type":None, "box":[0,0,original_image.size[0],original_image.size[1]]}
        for angle in response['results']:
            result = response['results'][angle]
            photo_score = 0
            document_score = 0
            document_type = None
            current_image = images_dict[angle]
            photo_box = [0, 0, current_image.size[0], current_image.size[1]]
            for detection in result:
                detection_data = result[detection]
                if detection == "id_photo":
                    photo_score = detection_data["score"]
                    photo_box = self.box_to_square(detection_data["original_box"])
                elif detection in ["id_card_front", "id_card_vertical", "passport"]:
                    if detection_data["score"] > document_score:
                        document_score = detection_data["score"]
                        document_type = detection

            if document_type == "id_card_vertical":
                # Vertical ID cards are detected assuming top of the document on the left. Rotate clockwise 90 deg to get right-side up.
                angle -= 90
                if angle < 0:
                    angle += 360
                rotated_box = [photo_box[1], original_image.size[0] - photo_box[2], photo_box[3], original_image.size[0]-photo_box[0]]
                photo_box = rotated_box


            overall_score = (document_score + photo_score)/2
            if overall_score > best_result["score"]:
                best_result = {"score":overall_score, "angle":angle, "document_type":document_type, "box":photo_box}
            if angle == 0:
                if overall_score > 0.95:
                    break
                else:
                    best_result["score"] = overall_score * 1.1

        return best_result

