import torch
import torch.nn.functional as F
from face_matcher.eval_transforms import ClassificationPresetEval
from face_matcher.image_processor import ImageProcessor
import numpy as np
import os, io
from PIL import Image

class FaceMatcher:
    def __init__(self, debug_mode=False) -> None:
        model_path = os.path.dirname(__file__) + "/facematch_model.pt"
        self.model = torch.load(model_path, map_location=torch.device("cpu"))
        self.transforms = ClassificationPresetEval()
        self.debug_mode = debug_mode
        self.image_processor = ImageProcessor(debug_mode=debug_mode)

    def get_vectors(self, images, face_detector_threshold=0.7):
        """Returns the embedding vectors for each input image

        Input: [List of image urls or image dictionaries]
        Output: [List of image dictionaries with embeddings under 'embedding' key]

        image dictionary format: {
            "filename": image_url,

            # PIL Image manipulation arguments. (not required)
            "prerotate_counterclockwise": rotation_deg,
            "prerotate_counterclockwise": rotation_deg,
            "crop": [left, upper, right, and lower pixel coordinates]

            # (Experimental) document detector. 
            # This will override PIL args with values from chahm's docv document detector
            "detect_document_photo": True

            # output populated by model
            "vector": vector embedding,
            "face_detected: True/False,
        }
        """
        self.set_face_detector_threshold(face_detector_threshold)
        images = self.image_processor(images)

        if len(images) > 0:
            batch = torch.stack([self.transforms(i["processed_image"]) for i in images])
            embeddings = self.model(batch)
            embeddings = F.normalize(embeddings.detach().cpu())

        for i, image in enumerate(images):
            image["vector"] = embeddings[i].tolist()
            del image["processed_image"]
            del image["raw_image"]

        return images

    def compare(self, a, b, return_vectors=False, face_detector_threshold=0.7):
        """Returns the cosine similarities between images in list a and images in list b. See input format below.

        Input: {a: [List of image dicts], b: [list of image dicts], return_vectors=True}. Dict must contain the keys "a" and "b".
        Output: [{a: image_url_a, b: image_url_b, similarity: cosine_similarity}]

        image dict format: {
            "filename": image_url,

            # PIL Image manipulation arguments. (not required)
            "prerotate_counterclockwise": rotation_deg,
            "postrotate_counterclockwise": rotation_deg,
            "crop": [left, upper, right, and lower pixel coordinates]
        }
        """
        input_imgs = a + b
        images_with_embeddings = self.get_vectors(
            input_imgs, face_detector_threshold=face_detector_threshold
        )
        a_embeds = images_with_embeddings[: len(a)]
        b_embeds = images_with_embeddings[len(a) :]

        out = []

        for a, b in zip(a_embeds, b_embeds):
            similarity = torch.nn.functional.cosine_similarity(
                torch.tensor(a["vector"]), torch.tensor(b["vector"]), dim=0
            )
            obj = {
                "a": a,
                "b": b,
                "similarity": float(similarity),
            }
            if not return_vectors:
                del obj["a"]["vector"]
                del obj["b"]["vector"]
            out.append(obj)
    
        return out
    
    def match_pair(self, id: bytearray, selfie: bytearray, return_vectors=False, face_detector_threshold=0.7):
        """Interface set up for docv team. Takes in a single id/selfie pair 

        Input: 
            id_bytearray: An ID image in bytearray format
            selfie_bytearray: A selfie image in bytearray format
            return_vectors: Whether to return vectors in metadata
            verbose: Whether to return vectors and detection metadata in response.
            face_detector_threshold: Sensitivity threshold for face detection, lower values make face detection more likely.
        Output: {final_score: id-selfie match score ranging from 0 to 1, id: id_metadata, selfie: selfie_metadata}
        """

        id = Image.open(io.BytesIO(id)).convert("RGB")
        selfie = Image.open(io.BytesIO(selfie)).convert("RGB")
        response = self.match(ids=[id], selfies=[selfie], return_vectors=return_vectors, face_detector_threshold=face_detector_threshold)
        
        return response[0]

    def match(self, ids, selfies, return_vectors=False, face_detector_threshold=0.7):
        """Returns the cosine similarities between a list of ids and selifes.

        Input: 
            ids: A list of PIL images of ids
            selfies: A list of PIL images of selfies
            return_vectors: Whether to return vectors in metadata
            face_detector_threshold: Sensitivity threshold for face detector. Lower to increase chances of face detection.
        Output: [{similarity: id-selfie facial similarity, id: id_metadata, selfie: selfie_metadata}, ...]
        """
        if len(ids) != len(selfies):
            raise f"Warning: The number of ids and selfies does not match. There are {len(ids)} ids and {len(selfies)} selfies."
        ids = [{"raw_image": im, "detect_document_photo": True} for im in ids]
        selfies = [{"raw_image": im} for im in selfies]

        comparisons = self.compare(a=ids, b=selfies, return_vectors=return_vectors, face_detector_threshold=face_detector_threshold)
        for comp in comparisons:
            comp["id"] = comp.pop("a")
            comp["selfie"] = comp.pop("b")

            # Test that all faces were detected
            crop = comp["id"]["crop"]
            docv_detected = crop[0] != 0
            doc_blazeface_detected = comp["id"]["face_detected"]
            selfie_detected = comp["selfie"]["face_detected"]
            faces_detected = selfie_detected and (docv_detected or doc_blazeface_detected)

            if faces_detected: # Fit score into standard vendor range
                comp["final_score"] = (comp["similarity"] / 2) + 0.5
            else:
                comp["final_score"] = None

        return comparisons
    
    def __call__(self, *args, **kwargs): # Can call face_matcher() or face_matcher.match()
        return self.match(*args, **kwargs)

    def set_face_detector_threshold(self, threshold: float):
        """
        Sets the face detector threshold to a given value.
        """
        # These parameters live pretty deep!
        self.image_processor.prp.face_detector.face_detector.score_threshold = threshold
        self.image_processor.prp.face_detector.face_detector.score_threshold_raw = np.log(
            threshold / (1 - threshold)
        )
