import numpy as np
from bodorientpy.homogenous.transformations import compose_matrix, \
    decompose_matrix
import pandas as pd
from matplotlib import pyplot as plt
from bodorientpy.markers.transformations import triangle2homogeous_transform
from bodorientpy.markers.triangle import Triangle
from bodorientpy.plot import draw_frame


apex0 = pd.Series(data=[0.,  0.25,  0.],
                  index=['x', 'y', 'z'])
apex1 = pd.Series(data=[0.5, -0.5,  0.],
                  index=['x', 'y', 'z'])
apex2 = pd.Series(data=[0.5,  0.5,  0.],
                  index=['x', 'y', 'z'])

triangle_orig = Triangle(apex0, apex1, apex2)
frame_orig = triangle2homogeous_transform(
    triangle_orig, triangle_mode='y-axis=1-2')
triangle = Triangle(apex0, apex1, apex2)

pos = np.array([1.0, 1.0, 1.0])
yaw = +1*np.pi/3
pitch = +1*np.pi/6
roll = -1*np.pi/6
axes = 'rzyx'

transform = compose_matrix(translate=pos,
                           angles=[yaw, pitch, roll],
                           axes=axes)
triangle.transform(transform)
frame = triangle2homogeous_transform(
    triangle, triangle_mode='y-axis=1-2')
_, _, angles, translate, _ = decompose_matrix(frame, axes=axes)

fig = plt.figure(figsize=(5, 5))
ax = fig.add_subplot(111, projection='3d')
draw_frame(ax=ax, frame=frame_orig)
triangle_orig.plot(ax=ax, apex_marker='wo')
draw_frame(ax=ax, frame=frame)
triangle.plot(ax=ax, apex_marker='wo')

celltext = [['{:0.2f}'.format(np.rad2deg(yaw)),
             '{:0.2f}'.format(np.rad2deg(angles[0]))],
            ['{:0.2f}'.format(np.rad2deg(pitch)),
             '{:0.2f}'.format(np.rad2deg(angles[1]))],
            ['{:0.2f}'.format(np.rad2deg(roll)),
             '{:0.2f}'.format(np.rad2deg(angles[2]))]]
ax.table(cellText=celltext,
         rowLabels=['yaw', 'pitch', 'roll'],
         colLabels=['theoric', 'estimate'],
         loc='bottom')

ax.set_xlim([-1, 1.5])
ax.set_ylim([-1, 1.5])
ax.set_zlim([-1, 1.5])
fig.show()