#! /usr/bin/env python3.4
#
#########################
#      $Author$
#      $Date$
#      $HeadURL$
#      $Revision$
#
#########################


import string
import re
import copy
import math
from pprint import pprint as pp
import numpy as np
import scipy
from scipy import interpolate
from PIL import Image
from scipy.misc import *
from enum import Enum
import os
from glob import glob


class Effect(Enum):
    rotate90 = 1
    rotate180 = 2
    rotate270 = 3
    flipHorizontally = 4
    flipVertically = 5
    transpose = 6

class Homography:
    def __init__(self, **kwargs):

        #homography given
        if "homographyMatrix" in kwargs:
            self.hom = kwargs["homographyMatrix"]
        #source and target points given (may be with effect)
        elif "sourcePoints" in kwargs and "targetPoints" in kwargs:
            if "effect" in kwargs:
                 if isinstance(kwargs["effect"], Effect) or kwargs["effect"] == None:
                     self.effect = kwargs["effect"]
                     self.hom = self.computeHomography(kwargs["sourcePoints"],kwargs["targetPoints"],effect=self.effect)
                 else:
                     raise TypeError("Effect must be instance of Effect or None")
            else:
                self.hom = self.computeHomography(kwargs["sourcePoints"],kwargs["targetPoints"])

        else:
            raise ValueError("homographyMatrix not passed into function")

        #ensure homography matrix obtained is good
        if len(self.hom) == 3:
            for item in self.hom:
                if len(item) == 3:
                    for item2 in item:
                        if type(item2) != type(1.0):

                            raise ValueError("Must be floats")
                else:
                    raise ValueError("Must be 3x3 matrix of floats")
        else:
            raise ValueError("Must be 3x3 matrix of floats")


        #Get inversion matrix
        self.homMat = np.array(self.hom)
        try:
            self.homMatInv = np.linalg.inv(self.homMat)
        except:
            pass



    def forwardProject(self, point):
        #Apply the homography to the given (x, y) point, and return the point (x , y ). This method takes in
        #and returns a two-element tuple.
        mat = np.array([[point[0]],[point[1]],[1]])
        res = self.homMat.dot(mat)
        resN1 = res[0]/res[2]
        resN2 = res[1]/res[2]

        return (np.round(resN1[0],2),np.round(resN2[0],2))


    def inverseProject(self, pointPrime):
        #Apply the inverse homography to the given (x , y ) point, and return the point (x, y). This method
        #takes in and returns a two element tuple.
        mat = np.array([pointPrime[0],pointPrime[1],1])
        res = self.homMatInv.dot(mat)

        return (np.round(res[0]/res[2],2),np.round(res[1]/res[2],2))

    def computeHomography(self, sourcePoints, targetPoints, effect=None):
        if len(sourcePoints) != 4 or len(targetPoints) != 4:
            raise ValueError("Length of points must be 4")
        elif type(sourcePoints) != type([(1,2),(1,2),(1,2),(1,2)]) or type(targetPoints) != type([(1,2),(1,2),(1,2),(1,2)]):
            raise ValueError("Wrong type")

        x1 = sourcePoints[0][0]
        x2 = sourcePoints[1][0]
        x3 = sourcePoints[2][0]
        x4 = sourcePoints[3][0]

        y1 = sourcePoints[0][1]
        y2 = sourcePoints[1][1]
        y3 = sourcePoints[2][1]
        y4 = sourcePoints[3][1]


        x11 = targetPoints[0][0]
        x21 = targetPoints[1][0]
        x31 = targetPoints[2][0]
        x41 = targetPoints[3][0]

        y11 = targetPoints[0][1]
        y21 = targetPoints[1][1]
        y31 = targetPoints[2][1]
        y41 = targetPoints[3][1]


        if effect == Effect.flipHorizontally:
            x11,y11,x21,y21,x31,y31,x41,y41 = x31,y31,x41,y41,x11,y11,x21,y21
        elif effect == Effect.flipVertically:
            x11,y11,x21,y21,x31,y31,x41,y41 = x21,y21,x11,y11,x41,y41,x31,y31
        elif effect == Effect.rotate180:
            x11,y11,x21,y21,x31,y31,x41,y41 = x41,y41,x31,y31,x21,y21,x11,y11
        elif effect == Effect.rotate270:
            x11,y11,x21,y21,x31,y31,x41,y41 = x31,y31,x11,y11,x41,y41,x21,y21
        elif effect == Effect.rotate90:
            x11,y11,x21,y21,x31,y31,x41,y41 = x21,y21,x41,y41,x11,y11,x31,y31
        elif effect == Effect.transpose:
            x11,y11,x21,y21,x31,y31,x41,y41 = x11,y11,x31,y31,x21,y21,x41,y41

        A = [[x1, y1, 1, 0, 0, 0, -x11*x1, -x11*y1], [0, 0, 0, x1, y1, 1, -y11*x1, -y11*y1]]
        A.append([x2, y2, 1, 0, 0, 0, -x21*x2, -x21*y2])
        A.append([0, 0, 0, x2, y2, 1, -y21*x2, -y21*y2])
        A.append([x3, y3, 1, 0, 0, 0, -x31*x3, -x31*y3])
        A.append([0, 0, 0, x3, y3, 1, -y31*x3, -y31*y3])
        A.append([x4, y4, 1, 0, 0, 0, -x41*x4, -x41*y4])
        A.append([0, 0, 0, x4, y4, 1, -y41*x4, -y41*y4])
        b = [[x11],[y11],[x21],[y21],[x31],[y31],[x41],[y41]]

        A = np.array(A)
        b = np.array(b)
        h = np.linalg.solve(A, b)

        return [[float(h[0][0]),float(h[1][0]),float(h[2][0])], [float(h[3][0]),float(h[4][0]),float(h[5][0])], [float(h[6][0]),float(h[7][0]),1.0]]


class Transformation:
    def __init__(self, sourceImage, homography=None):
        if type(homography) != type(Homography(homographyMatrix=[[2.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 1.0]])) and homography != None:
            raise TypeError("Second argument must be of type homography")
        if not isinstance(sourceImage,np.ndarray):
            raise TypeError("First argument must be of type ndarray")
        self.hom = homography
        self.srcImage = sourceImage


    def setupTransformation(self, targetPoints,effect=None):
        #Takes in targetPoints, a list of 4 elements, each of which is a point tuple. These points can be used
        #to identify the range of iteration in the target image
        if len(targetPoints) == 4:
            for item in targetPoints:
                if type(item) == type((1,2)):
                    pass
                else:
                    raise TypeError("Must be tuples in list")
        else:
            raise ValueError("Must have 4 items in list")

        xarray = [targetPoints[0][0], targetPoints[1][0], targetPoints[2][0], targetPoints[3][0]]
        self.maxX = max(xarray)
        self.minX = min(xarray)
        yarray = [targetPoints[0][1], targetPoints[1][1], targetPoints[2][1], targetPoints[3][1]]
        self.maxY = max(yarray)
        self.minY = min(yarray)

        yMax,xMax = self.srcImage.shape[0], self.srcImage.shape[1]
        src = [(0,0),(xMax-1,0),(0,yMax-1),(xMax-1,yMax-1)]

        if self.hom == None:
            self.hom = Homography(sourcePoints=src,targetPoints=targetPoints,effect=effect)

    def transformImage(self, containerImage):
        try:
            if not isinstance(containerImage,np.ndarray):
                raise TypeError("Argument must be of type ndarray")
            if not containerImage.ndim == 2:
                raise ValueError("Must have 2 dimensions")

            ylen, xlen = self.srcImage.shape
            xl, yl = np.arange(xlen),np.arange(ylen)
            sp = interpolate.RectBivariateSpline(yl,xl,self.srcImage,kx=1,ky=1)

            x = self.minX
            y = self.minY

            while y <= self.maxY:
                while x <= self.maxX:
                    xNew,yNew = self.hom.inverseProject((x,y))

                    if xNew <= xlen - 1 and yNew <= ylen - 1 and xNew >= 0 and yNew >= 0:
                        containerImage[y,x] = sp(yNew,xNew)

                    x += 1
                y += 1
                x = self.minX
            return containerImage
        except:
            if not isinstance(containerImage,np.ndarray):
                raise TypeError("Argument must be of type ndarray")
            if not containerImage.ndim == 3:
                raise ValueError("Must have 3 dimensions")

            ylen, xlen, color = self.srcImage.shape
            xl, yl = np.arange(xlen),np.arange(ylen)
            sp1 = interpolate.RectBivariateSpline(yl,xl,self.srcImage[...,0],kx=1,ky=1)
            sp2 = interpolate.RectBivariateSpline(yl,xl,self.srcImage[...,1],kx=1,ky=1)
            sp3 = interpolate.RectBivariateSpline(yl,xl,self.srcImage[...,2],kx=1,ky=1)

            x = self.minX
            y = self.minY

            while y <= self.maxY:
                while x <= self.maxX:
                    xNew,yNew = self.hom.inverseProject((x,y))

                    if xNew <= xlen - 1 and yNew <= ylen - 1 and xNew >= 0 and yNew >= 0:
                        containerImage[y,x,0] = sp1(yNew,xNew)
                        containerImage[y,x,1] = sp2(yNew,xNew)
                        containerImage[y,x,2] = sp3(yNew,xNew)
                    x += 1
                y += 1
                x = self.minX
            return containerImage


class ColorTransformation(Transformation):
    def __init__(self, sourceImage, homography=None):
        Transformation.__init__(self, sourceImage, homography)
        if not sourceImage.ndim == 3:
            raise ValueError("Must have 3 dimensions")


    def transformImage(self, containerImage):
        if not isinstance(containerImage,np.ndarray):
            raise TypeError("Argument must be of type ndarray")
        if not containerImage.ndim == 3:
            raise ValueError("Must have 3 dimensions")

        ylen, xlen, color = self.srcImage.shape
        xl, yl = np.arange(xlen),np.arange(ylen)
        sp1 = interpolate.RectBivariateSpline(yl,xl,self.srcImage[...,0],kx=1,ky=1)
        sp2 = interpolate.RectBivariateSpline(yl,xl,self.srcImage[...,1],kx=1,ky=1)
        sp3 = interpolate.RectBivariateSpline(yl,xl,self.srcImage[...,2],kx=1,ky=1)

        x = self.minX
        y = self.minY

        while y <= self.maxY:
            while x <= self.maxX:
                xNew,yNew = self.hom.inverseProject((x,y))

                if xNew <= xlen - 1 and yNew <= ylen - 1 and xNew >= 0 and yNew >= 0:
                    containerImage[y,x,0] = sp1(yNew,xNew)
                    containerImage[y,x,1] = sp2(yNew,xNew)
                    containerImage[y,x,2] = sp3(yNew,xNew)
                x += 1
            y += 1
            x = self.minX
        return containerImage


if __name__ == '__main__':
    print("In main")


    #source[...,0]
