#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
import scipy
from scipy.ndimage import map_coordinates

# this is code from here: http://www.lfd.uci.edu/~gohlke/code/transformations.py.html
from transformations import rotation_matrix, translation_matrix


# First in 2D:
# 
# Create a meshgrid for the coordinates, stack it to create "position vectors", transform them

# In[2]:


dim=100

img=np.zeros((dim,dim))
img[dim/5:dim/4,dim/4:dim*3/4]=1

get_ipython().run_line_magic('matplotlib', 'inline')
plt.figure("2d before trafo")
plt.imshow(img, interpolation="none")
plt.show()


# In[3]:


# create a meshgrid for the coordinates
line=np.arange(dim)
coords=np.meshgrid(line, line)
xy_coords=np.vstack([coords[0].reshape(-1), coords[1].reshape(-1)])

# create a transformation matrix
angle=30./180.*np.pi
c=np.cos(angle)
s=np.sin(angle)
mat=np.array([[c,s],[-s,c]])

# apply the transformation matrix
# please note: the coordinates are not homogeneous.
# for the 3D case, I've added code for homogeneous coordinates, you might want to look at that
# please also note: rotation is always around the origin:
# since I want the origin to be in the image center, I had to substract dim/2, rotate, then add it again
xy_coords=np.dot(mat, xy_coords-dim/2)+dim/2

# undo the stacking and reshaping
x=xy_coords[0,:]
y=xy_coords[1,:]
x=x.reshape((dim,dim))
y=y.reshape((dim,dim))

# create a list for map_coordinates
# ATTENTION: the coordinate system for map_coordinates seems messed up
# it is necessary to reorder the dimensions like this
# probably my fault, but I haven't figured out why, yet
new_coords=[y,x]

# use map_coordinates to sample values for the new image
new_img=scipy.ndimage.map_coordinates(img,new_coords, order=0)


# In[4]:


plt.figure("2d after trafo")
plt.imshow(new_img, interpolation="none")
plt.show()


# Now the 3D case

# In[5]:


dim=40

vol=np.zeros((dim,dim,dim))
vol[:,dim/5:dim/4,dim/4:dim*3/4]=1

plt.figure("2d slice of 3d volume, before trafo")
plt.imshow(vol[0,:], interpolation="none")
plt.show()


# In[6]:


ax=np.arange(dim)
coords=np.meshgrid(ax,ax,ax)

# stack the meshgrid to position vectors, center them around 0 by substracting dim/2
xyz=np.vstack([coords[0].reshape(-1)-float(dim)/2,     # x coordinate, centered
               coords[1].reshape(-1)-float(dim)/2,     # y coordinate, centered
               coords[2].reshape(-1)-float(dim)/2,     # z coordinate, centered
               np.ones((dim,dim,dim)).reshape(-1)])    # 1 for homogeneous coordinates

# create transformation matrix
mat=rotation_matrix(0.1*np.pi,(0,1,0))
#mat=translation_matrix(np.array([1,0,0]).transpose(1,0,2))

# apply transformation
transformed_xyz=np.dot(mat, xyz)

# extract coordinates, don't use transformed_xyz[3,:] that's the homogeneous coordinate, always 1
x=transformed_xyz[0,:]+float(dim)/2
y=transformed_xyz[1,:]+float(dim)/2
z=transformed_xyz[2,:]+float(dim)/2

x=x.reshape((dim,dim,dim))
y=y.reshape((dim,dim,dim))
z=z.reshape((dim,dim,dim))

# the coordinate system seems to be strange, it has to be ordered like this
new_xyz=[y,x,z]

# sample
new_vol=scipy.ndimage.map_coordinates(vol,new_xyz, order=0)


# In[7]:


plt.figure("2d slice of 3d volume, after trafo")
plt.imshow(new_vol[20,:,:], interpolation="none")
plt.show()


# In[ ]:





# In[ ]: