# (C) 2021: Hans Georg Schaathun <georg@schaathun.net>

import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def getVertices():
   points = [ [-1,0.5,0.5], [+1,0.5,0.5], [0,-0.5,0.5], [0,0.5,-0.5] ]
   vertices = []

   for i in range(4):
       face = points.copy()
       del face[i]
       vertices = vertices + face
   
   return np.array(vertices,dtype=float)

def getPoly(v):
   ob = Poly3DCollection(v, linewidths=1, alpha=0.2)
   ob.set_facecolor( [0.5, 0.5, 1] )
   ob.set_edgecolor([0,0,0])
   return ob

def getFig():
   plt.ion()
   fig = plt.figure()
   ax = fig.add_subplot(111, projection='3d')
   s = [-2,-2,-2,2,2,2]
   ax.auto_scale_xyz(s,s,s)
   plt.show()
   return fig,ax

def rotZ(theta):
    return np.array( [ [ np.cos(theta), -np.sin(theta), 0 ],
      [ np.sin(theta), np.cos(theta), 0 ],
      [ 0, 0, 1 ] ], dtype=float )

def toHom(v):
    xlen,ylen = v.shape
    vcol = np.zeros([xlen,1],dtype=float)+1
    return np.concatenate((v,vcol),axis=1)
def toHomMatrix(rotation,translation):
    print( rotation.shape )
    print( translation.shape )
    t = translation[:,np.newaxis]
    print( t.shape )
    m = np.concatenate((rotation,t),axis=1)
    r = np.array([0,0,0,1],dtype=float)
    r = r[np.newaxis,:]
    return np.concatenate((m,r),axis=0)

v = getVertices() 

translation = np.array([-1,0,1],dtype=float)
v2 = v + translation


v3 = v.dot(rotZ(np.pi/4))

ht = toHomMatrix(rotZ(np.pi/4),translation)

v4hom = ht.dot(toHom(v).transpose()).transpose()
v4 = v4hom[:,:-1]


def plotFigs(ax):
   ax.add_collection3d(getPoly(v))
   # ax.add_collection3d(getPoly(v2))
   # ax.add_collection3d(getPoly(v3))
   ax.add_collection3d(getPoly(v4))