import cv2
import numpy as np
from scipy.ndimage.filters import gaussian_filter

from .pose_utils import _get_keypoints, _pad_image
from insightface import model_zoo
from dofaker.utils import download_file, get_model_url
from dofaker.transforms import center_crop, pad


class PoseTransfer:

    def __init__(self,
                 name='pose_transfer',
                 root='weights/models',
                 pose_estimator=None):
        assert pose_estimator is not None, "The pose_estimator of PoseTransfer shouldn't be None"
        self.pose_estimator = pose_estimator
        _, model_file = download_file(get_model_url(name),
                                      save_dir=root,
                                      overwrite=False)
        providers = model_zoo.model_zoo.get_default_providers()
        self.session = model_zoo.model_zoo.PickableInferenceSession(
            model_file, providers=providers)

        self.input_mean = 127.5
        self.input_std = 127.5
        inputs = self.session.get_inputs()
        self.input_names = []
        for inp in inputs:
            self.input_names.append(inp.name)
        outputs = self.session.get_outputs()
        output_names = []
        for out in outputs:
            output_names.append(out.name)
        self.output_names = output_names
        assert len(
            self.output_names
        ) == 1, "The output number of PoseTransfer model should be 1, but got {}, please check your model.".format(
            len(self.output_names))
        output_shape = outputs[0].shape
        input_cfg = inputs[0]
        input_shape = input_cfg.shape
        self.input_shape = input_shape
        print('pose transfer shape:', self.input_shape)

    def forward(self, source_image, target_image, image_format='rgb'):
        h, w, c = source_image.shape
        if image_format == 'rgb':
            pass
        elif image_format == 'bgr':
            source_image = cv2.cvtColor(source_image, cv2.COLOR_BGR2RGB)
            target_image = cv2.cvtColor(target_image, cv2.COLOR_BGR2RGB)
            image_format = 'rgb'
        else:
            raise UserWarning(
                "PoseTransfer not support image format {}".format(image_format))
        imgA = self._resize_and_pad_image(source_image)
        kptA = self._estimate_keypoints(imgA, image_format=image_format)
        mapA = self._keypoints2heatmaps(kptA)

        imgB = self._resize_and_pad_image(target_image)
        kptB = self._estimate_keypoints(imgB)
        mapB = self._keypoints2heatmaps(kptB)

        imgA_t = (imgA.astype('float32') - self.input_mean) / self.input_std
        imgA_t = imgA_t.transpose([2, 0, 1])[None, ...]
        mapA_t = mapA.transpose([2, 0, 1])[None, ...]
        mapB_t = mapB.transpose([2, 0, 1])[None, ...]
        mapAB_t = np.concatenate((mapA_t, mapB_t), axis=1)
        pred = self.session.run(self.output_names, {
            self.input_names[0]: imgA_t,
            self.input_names[1]: mapAB_t
        })[0]
        target_image = pred.transpose((0, 2, 3, 1))[0]
        bgr_target_image = np.clip(
            self.input_std * target_image + self.input_mean, 0,
            255).astype(np.uint8)[:, :, ::-1]
        crop_size = (256,
                     min((256 * target_image.shape[1] // target_image.shape[0]),
                         176))
        bgr_image = center_crop(bgr_target_image, crop_size)
        bgr_image = cv2.resize(bgr_image, (w, h), interpolation=cv2.INTER_CUBIC)
        return bgr_image

    def get(self, source_image, target_image, image_format='rgb'):
        return self.forward(source_image, target_image, image_format)

    def _resize_and_pad_image(self, image: np.ndarray, size=256):
        w = size * image.shape[1] // image.shape[0]
        w_box = min(w, size * 11 // 16)
        image = cv2.resize(image, (w, size), interpolation=cv2.INTER_CUBIC)
        image = center_crop(image, (size, w_box))
        image = pad(image,
                    size - w_box,
                    size - w_box,
                    size - w_box,
                    size - w_box,
                    fill=255)
        image = center_crop(image, (size, size))
        return image

    def _estimate_keypoints(self, image: np.ndarray, image_format='rgb'):
        keypoints = self.pose_estimator.get(image, image_format)
        keypoints = keypoints[0] if len(keypoints) > 0 else np.zeros(
            (18, 3), dtype=np.int32)
        keypoints[np.where(keypoints[:, 2] == 0), :2] = -1
        keypoints = keypoints[:, :2]
        return keypoints

    def _keypoints2heatmaps(self, keypoints, size=256):
        heatmaps = np.zeros((size, size, keypoints.shape[0]), dtype=np.float32)
        for k in range(keypoints.shape[0]):
            x, y = keypoints[k]
            if x == -1 or y == -1:
                continue
            heatmaps[y, x, k] = 1.0
        return heatmaps