import torch
import torch.nn.functional as F
from PIL import Image
from facenet_pytorch import MTCNN
from transformers import Pipeline


class DeepFakePipeline(Pipeline):
    def __init__(self, **kwargs):
        Pipeline.__init__(self, **kwargs)

    def _sanitize_parameters(self, **kwargs):
        return {}, {}, {}

    def preprocess(self, inputs):
        return inputs

    def _forward(self, input):
        return input

    def postprocess(self, confidences):
        out = {"confidences": confidences}
        return out

    def predict(self, input_image: Image.Image):
        DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        mtcnn = MTCNN(
            select_largest=False,
            post_process=False,
            device=DEVICE)
        mtcnn.to(DEVICE)
        model = self.model.model
        model.to(DEVICE)

        face = mtcnn(input_image)
        if face is None:
            raise Exception('No face detected')

        face = face.unsqueeze(0)  # add the batch dimension
        face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
        face = face.to(DEVICE)
        face = face.to(torch.float32)
        face = face / 255.0

        with torch.no_grad():
            output = torch.sigmoid(model(face).squeeze(0))
            real_prediction = 1 - output.item()
            fake_prediction = output.item()
            confidences = {
                'real': real_prediction,
                'fake': fake_prediction
            }
        return self.postprocess(confidences)