import cv2
import numpy as np
import torch

from spiga.data.loaders.augmentors.modern_posit import PositPose
from spiga.data.loaders.augmentors.heatmaps import Heatmaps
from spiga.data.loaders.augmentors.boundary import AddBoundary
from spiga.data.loaders.augmentors.landmarks import HorizontalFlipAug, RSTAug, OcclusionAug, \
                                              LightingAug, BlurAug, TargetCropAug


def get_transformers(data_config):

    # Data augmentation
    aug_names = data_config.aug_names
    augmentors = []

    if 'flip' in aug_names:
        augmentors.append(HorizontalFlipAug(data_config.database.ldm_flip_order, data_config.hflip_prob))
    if 'rotate_scale' in aug_names:
        augmentors.append(RSTAug(data_config.angle_range, data_config.scale_min,
                                 data_config.scale_max, data_config.trl_ratio))
    if 'occlusion' in aug_names:
        augmentors.append(OcclusionAug(data_config.occluded_min_len,
                                       data_config.occluded_max_len,
                                       data_config.database.num_landmarks))
    if 'lighting' in aug_names:
        augmentors.append(LightingAug(data_config.hsv_range_min, data_config.hsv_range_max))
    if 'blur' in aug_names:
        augmentors.append(BlurAug(data_config.blur_prob, data_config.blur_kernel_range))

    # Crop mandatory
    augmentors.append(TargetCropAug(data_config.image_size, data_config.ftmap_size, data_config.target_dist))
    # Opencv style
    augmentors.append(ToOpencv())

    # Gaussian heatmaps
    if 'heatmaps2D' in aug_names:
        augmentors.append(Heatmaps(data_config.database.num_landmarks, data_config.ftmap_size,
                                   data_config.sigma2D, norm=data_config.heatmap2D_norm))

    if 'boundaries' in aug_names:
        augmentors.append(AddBoundary(num_landmarks=data_config.database.num_landmarks,
                                      map_size=data_config.ftmap_size,
                                      sigma=data_config.sigmaBD))
    # Pose generator
    if data_config.generate_pose:
        augmentors.append(PositPose(data_config.database.ldm_ids,
                                    focal_ratio=data_config.focal_ratio,
                                    selected_ids=data_config.posit_ids,
                                    max_iter=data_config.posit_max_iter))

    return augmentors


class ToOpencv:
    def __call__(self, sample):
        # Convert in a numpy array and change to GBR
        image = np.array(sample['image'])
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        sample['image'] = image
        return sample


class TargetCrop(TargetCropAug):
    def __init__(self, crop_size=256, target_dist=1.6):
        super(TargetCrop, self).__init__(crop_size, crop_size, target_dist)


class AddModel3D(PositPose):
    def __init__(self, ldm_ids, ftmap_size=(256, 256), focal_ratio=1.5, totensor=False):
        super(AddModel3D, self).__init__(ldm_ids, focal_ratio=focal_ratio)
        img_bbox = [0, 0, ftmap_size[1], ftmap_size[0]]  # Shapes given are inverted (y,x)
        self.cam_matrix = self._camera_matrix(img_bbox)

        if totensor:
            self.cam_matrix = torch.tensor(self.cam_matrix, dtype=torch.float)
            self.model3d_world = torch.tensor(self.model3d_world, dtype=torch.float)

    def __call__(self, sample={}):
        # Save intrinsic matrix and 3D model landmarks
        sample['cam_matrix'] = self.cam_matrix
        sample['model3d'] = self.model3d_world
        return sample