'''

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

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


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.1, 0.4), rotation_amount=15, warp_amount=0.05):
        self.canvas_size = canvas_size
        self.scale_amount = scale_amount
        self.rotation_amount = rotation_amount
        self.warp_amount = warp_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_homography2(self):
        # assume that the rectangle images have been rescaled to the fit the canvas
        final_transform = np.eye(3, dtype=np.float64)

        # translate the cornerpoints such that the centre of the card is at the origin
        tx = self.canvas_size[0] / 2.
        ty = self.canvas_size[1] / 2.
        translation_matrix = np.array([[1, 0, -tx],
                                       [0, 1, -ty],
                                       [0, 0, 1]], dtype=np.float64)
        final_transform = np.dot(translation_matrix, final_transform)

        # apply a random rescale (to produce image samples with rectangles 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.]], dtype=np.float64)
        final_transform = np.dot(zoom_matrix, final_transform)

        # 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]], dtype=np.float64)
        final_transform = np.dot(rotation_matrix, final_transform)

        # generate a small amount of perspective warping
        A = random.uniform(-self.warp_amount, self.warp_amount)
        B = random.uniform(-self.warp_amount, self.warp_amount)
        warp_matrix = np.array([[1, 0., 0.],
                                [0., 1, 0.],
                                [A, B, 1.]], dtype=np.float64)
        final_transform = np.dot(warp_matrix, final_transform)

        # translate back to centre of canvas
        translation_matrix2 = np.array([[1, 0, tx],
                                       [0, 1, ty],
                                       [0, 0, 1]], dtype=np.float64)
        final_transform = np.dot(translation_matrix2, final_transform)

        return final_transform

    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
        homography = self._random_homography2()
        # only interested in the 8 degrees of freedom, as a tuple
        h = tuple(homography.flatten()[:8])

        # 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 rectangle 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)

        cornerpoints_original = np.array([[0, 0],
                                          [self.canvas_size[0], 0],
                                          [self.canvas_size[0], self.canvas_size[1]],
                                          [0, self.canvas_size[1]]], dtype=np.float64)

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

        cornerpoints_warped = self._transform_cornerpoints(np.linalg.inv(homography), cornerpoints_original)

        for j in range(4):
            cornerpoints_warped = cornerpoints_warped.astype(int)
            cv2.circle(img_background_new, tuple(cornerpoints_warped[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 = (1, 2)
    WARP_AMOUNT = 0.00005

    photographer = PhotoSynthesizer(CANVAS_SIZE, scale_amount=SCALE_AMOUNT, warp_amount=WARP_AMOUNT)


    while True:
        photographer.snap(IMG_RECTANGLE, IMG_BACKGROUND)