'''

USAGE: $python this_script.py ipod.jpg background.jpg

'''
import random
from PIL import Image
import cv2
import numpy as np


class BBox(object):
    '''
    Given a set of cornerpoints (list of 4 2D position vectors), computes a tight bounding box around them
    '''

    def __init__(self, cornerpoints):
        self.left = np.min(cornerpoints[:, 0])
        self.right = np.max(cornerpoints[:, 0])
        self.upper = np.max(cornerpoints[:, 1])
        self.lower = np.min(cornerpoints[:, 1])

        self.width = self.right - self.left
        self.height = self.lower - self.upper

        self.coords = np.array([[self.left, self.upper],
                                [self.right, self.upper],
                                [self.right, self.lower],
                                [self.left, self.lower]])


class PhotoSynthesizer(object):
    '''
    Class for synthesizing "photos" of rectangles. Simulates perspective distortion, scale, rotation and translation.

    Convention for box/cornerpoints ordering:

            [top-left, top-right, bottom-right, bottom-left]

    i.e.    clockwise from top-left.
    '''
    def __init__(self, canvas_size,
                 scale_amount=(0.3, 0.7), rotation_amount=15, distortion_amount=50):

        self.canvas_size = canvas_size

        self.scale_amount = scale_amount
        self.rotation_amount = rotation_amount
        self.distortion_amount = distortion_amount

    def _to_homogenous(self, nparr):
        homogenous = np.c_[nparr, np.ones(4)]
        return homogenous

    def _from_homogenous(self, nparr):
        outarr = np.zeros((4, 2))
        outarr[:, 0] = nparr[:, 0] / nparr[:, 2]
        outarr[:, 1] = nparr[:, 1] / nparr[:, 2]
        return outarr

    def _transform_cornerpoints(self, matrix, coordinate_list):
        # turn the coords array into homogeneous coordinates, to do matrix operations on them
        coords_homogeneous = self._to_homogenous(coordinate_list)
        # transpose the cornerpoints arrays so that they can be multiplied by matrices
        coords_homogeneous = np.transpose(coords_homogeneous)

        # perform dot
        coords_homogeneous = np.dot(matrix, coords_homogeneous)

        # un-transpose
        coords_homogeneous = np.transpose(coords_homogeneous)
        # convert from homogeneous to 2D
        return self._from_homogenous(coords_homogeneous)

    def _random_homography(self):

        # assume that the ID card images have been rescaled to the fit the canvas
        w_original = self.canvas_size[0]
        h_original = self.canvas_size[1]

        cornerpoints_original = np.array([[0, 0],
                                          [w_original, 0],
                                          [w_original, h_original],
                                          [0, h_original]], dtype=np.float32)

        # flip the y-axis
        cornerpoints = np.copy(cornerpoints_original)
        cornerpoints[:, 1] *= -1

        # translate the cornerpoints such that the centre of the card is at the origin
        translate_vector = np.array([w_original / 2, h_original / -2])
        cornerpoints = cornerpoints - translate_vector

        # perspective distort the cornerpoints
        jitter = np.random.uniform(-self.distortion_amount, self.distortion_amount, (4, 2))
        cornerpoints = (cornerpoints + jitter).astype(np.float32)

        # perform a random rotation on the card cornerpoints
        theta = np.pi / 180 * np.random.uniform(-self.rotation_amount, self.rotation_amount)
        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
                                    [np.sin(theta), np.cos(theta),  0],
                                    [0,                         0,  1]])
        cornerpoints = self._transform_cornerpoints(rotation_matrix, cornerpoints).astype(np.float32)

        # apply a random rescale (to produce image samples with IDs of varying sizes)
        # that preserves aspect ratio
        scale_factor = random.uniform(self.scale_amount[0], self.scale_amount[1])
        zoom_matrix = np.array([[scale_factor, 0., 0.],
                                [0., scale_factor, 0.],
                                [0., 0., 1.]])
        cornerpoints = self._transform_cornerpoints(zoom_matrix, cornerpoints)

        # work out the tightly enclosing bounding box around cornerpoints
        bbox = BBox(cornerpoints)
        # print('tight bounding box: ')
        # print(bbox.upper, bbox.left, bbox.lower, bbox.right)

        # work out the maximum amounts that the card can travel horizontally and vertically, while remaining
        # relatively non-truncated
        upperbound = (self.canvas_size[1] / 2) + bbox.height / 10.
        lowerbound = -upperbound
        rightbound = (self.canvas_size[0] / 2) + bbox.width / 10.
        leftbound = -rightbound

        # print('upper, left, lower, right boundaries:')
        # print(upperbound, leftbound, lowerbound, rightbound)

        # these are the maximum amounts the card can travel in either direction before becoming *too* truncated
        max_travel_right = rightbound - bbox.right
        max_travel_left = bbox.left - leftbound
        max_travel_up = upperbound - bbox.upper
        max_travel_down = bbox.lower - lowerbound

        # print('max travel upper, left, lower, right')
        # print(max_travel_up, max_travel_left, max_travel_down, max_travel_right)

        random_translate = np.array([random.uniform(-max_travel_left, max_travel_right),
                                     random.uniform(-max_travel_down, max_travel_up)])

        # print('random translation vector:')
        # print(random_translate)

        cornerpoints = cornerpoints + random_translate

        # # translate the cornerpoints centre back from the cartesian origin to the image centre
        cornerpoints = (cornerpoints + translate_vector).astype(np.float32)

        # flip the y-axis back
        cornerpoints[:, 1] *= -1

        H = cv2.getPerspectiveTransform(cornerpoints_original, cornerpoints)
        # PIL is weird and insists on using the INVERSE homography matrix
        J = np.linalg.inv(H)
        # only interested in the 8 degrees of freedom, as a tuple
        h = tuple(J.flatten()[:8])

        return h, cornerpoints

    def snap(self, img_rectangle, img_background):
        '''
        :param img_rectangle: PIL image of rectangle
        :param img_background_new: PIL image of background
        :return: cv2 image - composite of rectangle and background with random perspective
        '''

        img_rectangle_new = img_rectangle.convert("RGBA")
        img_rectangle_new = img_rectangle_new.resize(self.canvas_size)

        # compute a random projective transform homography
        h, cornerpoints_new = self._random_homography()

        # apply the transform
        new_image = img_rectangle_new.transform(self.canvas_size, Image.PERSPECTIVE, h, Image.BICUBIC)

        # rescale the background to fit the canvas
        img_background_new = img_background.resize(self.canvas_size)

        # paste projective-transformed ID onto background
        img_background_new.paste(new_image, (0, 0), new_image)

        # convert the PIL image to opencv:
        img_background_new = np.array(img_background_new)
        img_background_new = cv2.cvtColor(img_background_new, cv2.COLOR_RGB2BGR)

        cols = [(225, 0, 0), (0, 225, 0), (0, 0, 225), (225, 225, 0)]

        for j in range(4):
            cornerpoints_new = cornerpoints_new.astype(int)
            cv2.circle(img_background_new, tuple(cornerpoints_new[j]), 5, color=cols[j], thickness=-1)

        show(img_background_new)

def show(im, winname=''):
    cv2.imshow(winname, im)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    import sys

    IMG_RECTANGLE = Image.open(sys.argv[1])
    IMG_BACKGROUND = Image.open(sys.argv[2])
    CANVAS_SIZE = (640, 410)
    SCALE_AMOUNT = (0.4, 1.1)
    DISTORTION_AMOUNT = 50

    photographer = PhotoSynthesizer(CANVAS_SIZE)

    while True:
        photographer.snap(IMG_RECTANGLE, IMG_BACKGROUND)