import numpy as np
from math import cos, sin, pi
import cv2
import random

from config import config, TransformationParams
from data_prep import map_coco_to_personlab

class AugmentSelection:

    def __init__(self, flip=False, degree = 0., crop = (0,0), scale = 1.):
        self.flip = flip
        self.degree = degree #rotate
        self.crop = crop #shift actually
        self.scale = scale

    @staticmethod
    def random():
        flip = random.uniform(0.,1.) > TransformationParams.flip_prob
        degree = random.uniform(-1.,1.) * TransformationParams.max_rotate_degree
        scale = (TransformationParams.scale_max - TransformationParams.scale_min)*random.uniform(0.,1.)+TransformationParams.scale_min \
            if random.uniform(0.,1.) < TransformationParams.scale_prob else 1.
        x_offset = int(random.uniform(-1.,1.) * TransformationParams.center_perterb_max);
        y_offset = int(random.uniform(-1.,1.) * TransformationParams.center_perterb_max);

        return AugmentSelection(flip, degree, (x_offset,y_offset), scale)

    @staticmethod
    def unrandom():
        flip = False
        degree = 0.
        scale = 1.
        x_offset = 0
        y_offset = 0

        return AugmentSelection(flip, degree, (x_offset,y_offset), scale)

    def affine(self, center=(config.IMAGE_SHAPE[1]//2, config.IMAGE_SHAPE[0]//2) , scale_self=1.):

        # the main idea: we will do all image transformations with one affine matrix.
        # this saves lot of cpu and make code significantly shorter
        # same affine matrix could be used to transform joint coordinates afterwards

        A = self.scale * cos(self.degree / 180. * pi )
        B = self.scale * sin(self.degree / 180. * pi )

        # scale_size = TransformationParams.target_dist / scale_self * self.scale
        scale_size = TransformationParams.target_dist / self.scale

        (width, height) = center
        center_x = width + self.crop[0]
        center_y = height + self.crop[1]

        center2zero = np.array( [[ 1., 0., -center_x],
                                 [ 0., 1., -center_y ],
                                 [ 0., 0., 1. ]] )

        rotate = np.array( [[ A, B, 0 ],
                           [ -B, A, 0 ],
                           [  0, 0, 1. ] ])

        scale = np.array( [[ scale_size, 0, 0 ],
                           [ 0, scale_size, 0 ],
                           [  0, 0, 1. ] ])

        flip = np.array( [[ -1 if self.flip else 1., 0., 0. ],
                          [ 0., 1., 0. ],
                          [ 0., 0., 1. ]] )

        center2center = np.array( [[ 1., 0., config.IMAGE_SHAPE[1]//2],
                                   [ 0., 1., config.IMAGE_SHAPE[0]//2 ],
                                   [ 0., 0., 1. ]] )

        # order of combination is reversed
        combined = center2center.dot(flip).dot(scale).dot(rotate).dot(center2zero)

        return combined[0:2]

class Transformer:

    @staticmethod
    def transform(img, masks, keypoints, aug=AugmentSelection.random()):

        # warp picture and mask
        M = aug.affine(center=(img.shape[1]//2, img.shape[0]//2))
        cv_shape = (config.IMAGE_SHAPE[1], config.IMAGE_SHAPE[0])

        # TODO: need to understand this, scale_provided[0] is height of main person divided by 368, caclulated in generate_hdf5.py
        # print(img.shape)
        # for i, img in enumerate(input_transform_targets):
        img = cv2.warpAffine(img, M, cv_shape, flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=(127,127,127))
        
        # concat = np.stack(output_transform_targets, axis=-1)
        # fix from https://github.com/octiapp/KerasPersonLab/issues/2
        # masks = cv2.warpAffine(masks, M, cv_shape, flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
        out_masks = np.zeros(cv_shape[::-1]+(masks.shape[-1],))
        for i in range(masks.shape[-1]):
            out_masks[:,:,i] = cv2.warpAffine(masks[:,:,i], M, cv_shape, flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
        masks = out_masks

        # warp key points
        #TODO: joint could be cropped by augmentation, in this case we should mark it as invisible.
        #update: may be we don't need it actually, original code removed part sliced more than half totally, may be we should keep it
        keypoints = map_coco_to_personlab(keypoints)
        original_points = keypoints.copy()
        # print keypoints
        original_points[:,:,2]=1  # we reuse 3rd column in completely different way here, it is hack
        converted_points = np.matmul(M, original_points.transpose([0,2,1])).transpose([0,2,1])
        keypoints[:,:,0:2]=converted_points

        cropped_kp = keypoints[:,:,0] >= config.IMAGE_SHAPE[1]
        cropped_kp = np.logical_or(cropped_kp, keypoints[:,:,1] >= config.IMAGE_SHAPE[0])
        cropped_kp = np.logical_or(cropped_kp, keypoints[:,:,0] < 0)
        cropped_kp = np.logical_or(cropped_kp, keypoints[:,:,1] < 0)

        keypoints[cropped_kp, 2] = 0
        

        # we just made image flip, i.e. right leg just became left leg, and vice versa

        if aug.flip:
            tmpLeft = keypoints[:, config.LEFT_KP, :]
            tmpRight = keypoints[:, config.RIGHT_KP, :]
            keypoints[:, config.LEFT_KP, :] = tmpRight
            keypoints[:, config.RIGHT_KP, :] = tmpLeft

        # print keypoints
        return img, masks, keypoints


[((54, 22, 56, 50), 'numpy.array', 'np.array', ({(54, 32, 56, 48): '[[1.0, 0.0, -center_x], [0.0, 1.0, -center_y], [0.0, 0.0, 1.0]]'}, {}), '([[1.0, 0.0, -center_x], [0.0, 1.0, -center_y], [0.0, 0.0, 1.0]])', True, 'import numpy as np\n'), ((58, 17, 60, 43), 'numpy.array', 'np.array', ({(58, 27, 60, 42): '[[A, B, 0], [-B, A, 0], [0, 0, 1.0]]'}, {}), '([[A, B, 0], [-B, A, 0], [0, 0, 1.0]])', True, 'import numpy as np\n'), ((62, 16, 64, 43), 'numpy.array', 'np.array', ({(62, 26, 64, 42): '[[scale_size, 0, 0], [0, scale_size, 0], [0, 0, 1.0]]'}, {}), '([[scale_size, 0, 0], [0, scale_size, 0], [0, 0, 1.0]])', True, 'import numpy as np\n'), ((66, 15, 68, 43), 'numpy.array', 'np.array', ({(66, 25, 68, 41): '[[-1 if self.flip else 1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]'}, {}), '([[-1 if self.flip else 1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0,\n    1.0]])', True, 'import numpy as np\n'), ((70, 24, 72, 52), 'numpy.array', 'np.array', ({(70, 34, 72, 50): '[[1.0, 0.0, config.IMAGE_SHAPE[1] // 2], [0.0, 1.0, config.IMAGE_SHAPE[0] //\n    2], [0.0, 0.0, 1.0]]'}, {}), '([[1.0, 0.0, config.IMAGE_SHAPE[1] // 2], [0.0, 1.0, config.\n    IMAGE_SHAPE[0] // 2], [0.0, 0.0, 1.0]])', True, 'import numpy as np\n'), ((91, 14, 91, 128), 'cv2.warpAffine', 'cv2.warpAffine', (), '', False, 'import cv2\n'), ((96, 20, 96, 63), 'numpy.zeros', 'np.zeros', ({(96, 29, 96, 62): 'cv_shape[::-1] + (masks.shape[-1],)'}, {}), '(cv_shape[::-1] + (masks.shape[-1],))', True, 'import numpy as np\n'), ((104, 20, 104, 52), 'data_prep.map_coco_to_personlab', 'map_coco_to_personlab', ({(104, 42, 104, 51): 'keypoints'}, {}), '(keypoints)', False, 'from data_prep import map_coco_to_personlab\n'), ((112, 21, 112, 89), 'numpy.logical_or', 'np.logical_or', ({(112, 35, 112, 45): 'cropped_kp', (112, 47, 112, 88): 'keypoints[:, :, (1)] >= config.IMAGE_SHAPE[0]'}, {}), '(cropped_kp, keypoints[:, :, (1)] >= config.IMAGE_SHAPE[0])', True, 'import numpy as np\n'), ((113, 21, 113, 68), 'numpy.logical_or', 'np.logical_or', ({(113, 35, 113, 45): 'cropped_kp', (113, 47, 113, 67): 'keypoints[:, :, (0)] < 0'}, {}), '(cropped_kp, keypoints[:, :, (0)] < 0)', True, 'import numpy as np\n'), ((114, 21, 114, 68), 'numpy.logical_or', 'np.logical_or', ({(114, 35, 114, 45): 'cropped_kp', (114, 47, 114, 67): 'keypoints[:, :, (1)] < 0'}, {}), '(cropped_kp, keypoints[:, :, (1)] < 0)', True, 'import numpy as np\n'), ((19, 15, 19, 36), 'random.uniform', 'random.uniform', ({(19, 30, 19, 32): '(0.0)', (19, 33, 19, 35): '(1.0)'}, {}), '(0.0, 1.0)', False, 'import random\n'), ((20, 17, 20, 39), 'random.uniform', 'random.uniform', ({(20, 32, 20, 35): '(-1.0)', (20, 36, 20, 38): '(1.0)'}, {}), '(-1.0, 1.0)', False, 'import random\n'), ((44, 25, 44, 54), 'math.cos', 'cos', ({(44, 29, 44, 52): '(self.degree / 180.0 * pi)'}, {}), '(self.degree / 180.0 * pi)', False, 'from math import cos, sin, pi\n'), ((45, 25, 45, 54), 'math.sin', 'sin', ({(45, 29, 45, 52): '(self.degree / 180.0 * pi)'}, {}), '(self.degree / 180.0 * pi)', False, 'from math import cos, sin, pi\n'), ((98, 31, 98, 142), 'cv2.warpAffine', 'cv2.warpAffine', (), '', False, 'import cv2\n'), ((22, 15, 22, 36), 'random.uniform', 'random.uniform', ({(22, 30, 22, 32): '(0.0)', (22, 33, 22, 35): '(1.0)'}, {}), '(0.0, 1.0)', False, 'import random\n'), ((23, 23, 23, 45), 'random.uniform', 'random.uniform', ({(23, 38, 23, 41): '-1.0', (23, 42, 23, 44): '1.0'}, {}), '(-1.0, 1.0)', False, 'import random\n'), ((24, 23, 24, 45), 'random.uniform', 'random.uniform', ({(24, 38, 24, 41): '-1.0', (24, 42, 24, 44): '1.0'}, {}), '(-1.0, 1.0)', False, 'import random\n'), ((21, 82, 21, 103), 'random.uniform', 'random.uniform', ({(21, 97, 21, 99): '(0.0)', (21, 100, 21, 102): '(1.0)'}, {}), '(0.0, 1.0)', False, 'import random\n')]