import os
import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt

from menpo.shape import PointCloud
import menpo.io as mio
from menpofit.transform import DifferentiableThinPlateSplines

"""
Explanation: Derivatives of a TPS
End of explanation
"""


src_landmarks = PointCloud(np.array([[-1, -1],
                                     [-1,  1],
                                     [ 1, -1],
                                     [ 1,  1]]))

tgt_landmarks = PointCloud(np.array([[-1, -1],
                                     [-1,  1],
                                     [ 1, -1],
                                     [ 1,  1]]))

"""
Explanation: We start by defining the source and target landmarks. Notice that, in this first example source = target!!!
End of explanation
"""


tps = DifferentiableThinPlateSplines(src_landmarks, tgt_landmarks)
np.allclose(tps.apply(src_landmarks).points, tgt_landmarks.points)

"""
Explanation: The warp can be effectively computed, although the rendering will not appear to be correct...
End of explanation
"""


x = np.arange(-1, 1, 0.01)
y = np.arange(-1, 1, 0.01)
xx, yy = np.meshgrid(x, y)
points = np.array([xx.flatten(1), yy.flatten(1)]).T

"""
Explanation: The next step is to define the set of points at which the derivative of the previous TPS warp must be evaluated. In this case, we use the function meshgrid to generate points inside the convex hull defined by the source landmarks.
End of explanation
"""


%matplotlib inline
dW_dxy = tps.d_dl(points)
reshaped = dW_dxy.reshape(xx.shape + (4,2))

#dW_dx
plt.subplot(241)
plt.imshow(reshaped[:,:,0,0])
plt.subplot(242)
plt.imshow(reshaped[:,:,1,0])
plt.subplot(243)
plt.imshow(reshaped[:,:,2,0])
plt.subplot(244)
plt.imshow(reshaped[:,:,3,0])

#dW_dy
plt.subplot(245)
plt.imshow(reshaped[:,:,0,1])
plt.subplot(246)
plt.imshow(reshaped[:,:,1,1])
plt.subplot(247)
plt.imshow(reshaped[:,:,2,1])
plt.subplot(248)
plt.imshow(reshaped[:,:,3,1])

"""
Explanation: We evaluate the derivative, reshape the output, and visualize the result.
End of explanation
"""


print(reshaped[1:5,1:5,0,0])
print(reshaped[1:5,1:5,0,1])

"""
Explanation: If everything goes as expected, the upper corner of the images defining the derivative of the warp wrt the x and y coordinates of the first of the source landmarks should both contain values close to 1.
End of explanation
"""


summed_x = np.sum(reshaped[:,:,:,0], axis=-1)
np.allclose(np.ones(xx.shape), summed_x)

plt.imshow(summed_x)

"""
Explanation: The sum of all the derivatives wrt the x coordinates should produce an all 1 image
End of explanation
"""


summed_y = np.sum(reshaped[:,:,:,1], axis=-1)
np.allclose(np.ones(xx.shape), summed_y)

plt.imshow(summed_y)

"""
Explanation: and so should the sum of all derivatives wrt the y coordinates.
End of explanation
"""


np.allclose(reshaped[:,:,:,0], reshaped[:,:,:,1])

"""
Explanation: Finally, the derivatives with respect to the x and y coordinates should be in this case exactly the same!!!
End of explanation
"""


