import torch
import torchvision.transforms as transforms

import cv2
import face_alignment
import numpy as np
from PIL import Image

from data_preprocessing.lm_based_image_align import LandmarkBasedImageAlign
from data_preprocessing.pose_estimation import PoseEstimator
from data_preprocessing.crop_image import ImageCropper
from data_preprocessing.foreground_extractor import ForegroundExtractor


class DataPreprocessor:
    def __init__(self, device, crop_smooth_alpha=0.9):
        self.device = device
        self.crop_smooth_alpha = crop_smooth_alpha

        self.face_alignment = LandmarkBasedImageAlign(output_size=1024, transform_size=1024)
        self.pose_estimator = PoseEstimator(device)
        self.cropper = ImageCropper()
        self.foreground_extractor = ForegroundExtractor(device)
        self.fa = face_alignment.FaceAlignment(
            face_alignment.LandmarksType.TWO_D, flip_input=False, device=device
        )

        self.transform = transforms.ToTensor()

        self.lm_woAlign_ema = None
        self.aligned_lm_ema = None

    def __call__(self, img, keep_bg=False):
        lm = self.fa.get_landmarks(np.array(img))

        if lm is None:
            detected_face = [0, 0, img.size[0], img.size[1]]
            lm = self.fa.get_landmarks(img, detected_faces=[detected_face])[0]
        else:
            lm = lm[0]
        if self.lm_woAlign_ema is None:
            self.lm_woAlign_ema = np.array(lm)
        self.lm_woAlign_ema = (self.lm_woAlign_ema * 0.9 + np.array(lm) * 0.1)
        img_woAlign, lm_woAlign = self.cropper(img, self.lm_woAlign_ema)
        lm_woAlign = np.array(lm_woAlign)
        img_woAlign = np.array(img_woAlign)

        img, aligned_lm = self.face_alignment(img, lm)
        intrinsics, pose = self.pose_estimator.predict_pose(img, aligned_lm)
        img, aligned_lm = self.cropper(img, aligned_lm)
        img = np.array(img)

        if not keep_bg:
            matte = self.foreground_extractor(img)
            img = (img * matte).astype(np.uint8)

        if self.aligned_lm_ema is None:
            self.aligned_lm_ema = aligned_lm
        self.aligned_lm_ema = 0.8 * self.aligned_lm_ema + aligned_lm * 0.2
        crop_params = cv2.estimateAffinePartial2D(self.aligned_lm_ema, lm_woAlign)[0]

        return img, img_woAlign, intrinsics, pose, crop_params

    def from_path(self, image_path, device, keep_bg=False):
        img = cv2.imread(image_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(img)

        img, _, intrinsics, pose, crop_params = self(img, keep_bg=keep_bg)

        img = np.transpose(img, (2, 0, 1))[None, :, :, :] / 255.
        img = (img * 2 - 1)
        img = torch.from_numpy(img).float()

        pose = torch.from_numpy(pose).unsqueeze(0).float()
        intrinsics = torch.from_numpy(intrinsics).unsqueeze(0).float()

        return {
            'image': img.to(device),
            'cam2world': pose.to(device),
            'intrinsics': intrinsics.to(device),
            'crop_params': crop_params
        }

    @staticmethod
    def realign(img, T):
        out = cv2.warpAffine(img, T, (img.shape[1], img.shape[0]))
        return out