import math

import matplotlib.pyplot as plt
from matplotlib.patches import Arc, FancyArrowPatch

"""
Draws a parameterized top view of a robot. 

This script is slightly over-engineered however it gives a nice and good display on how the robots look like from above,
especially how the wheels stand and the external (x,y) and internal (u,v) directions. 

Author: Tom Meulenkamp (t.meulenkamp@roboteamtwente.nl)
"""

R = 0.081                              # Robot radius [m]
Rw = 0.028                             # Wheel radius [m]
phi_deg = 30                           # Angle of the front wheels [deg]
theta_deg = 60                         # Angle of the back wheels  [deg]
phi_rad = phi_deg / 180 * math.pi      # Angle of front wheels     [rad]
theta_rad = theta_deg / 180 * math.pi  # Angle of the back wheels  [rad]
line_width_wheel = 10                  # The thickness of a wheel (not realistic)
annotation_offset = 0.005              # The offset that an annotation should have in comparison to its anchor point
RTT_purple = "#5a1c74"                 # RoboTeam purple for the accent colors

fig, ax = plt.subplots(figsize=(10, 10), dpi=100)

# Draws the robot
arc = Arc((0, 0), width=R * 2, height=R * 2, theta1=0, theta2=300, angle=120)
ax.plot([-math.sin(30/180*math.pi) * R, math.sin(30/180*math.pi) * R],
        [math.cos(30/180*math.pi) * R, math.cos(30/180*math.pi) * R], color="black")

# Lines going towards the wheels
wheels = [(math.cos(theta_rad) * R, -math.sin(theta_rad) * R),
          (-math.cos(theta_rad) * R, -math.sin(theta_rad) * R),
          (-math.cos(phi_rad) * R, math.sin(phi_rad) * R),
          (math.cos(phi_rad) * R, math.sin(phi_rad) * R)]
ax.plot([0, wheels[0][0]], [0, wheels[0][1]], color="grey")
ax.plot([0, wheels[1][0]], [0, wheels[1][1]], color="grey")
ax.plot([0, wheels[2][0]], [0, wheels[2][1]], color="grey")
ax.plot([0, wheels[3][0]], [0, wheels[3][1]], color="grey")
ax.text(0.01, wheels[3][1], "$R_r$ = " + str(R), fontsize=15)

# Wheels
ax.plot([wheels[0][0] - 0.5 * Rw * math.sin(theta_rad), wheels[0][0] + 0.5 * Rw * math.sin(theta_rad)],
        [wheels[0][1] - 0.5 * Rw * math.cos(theta_rad), wheels[0][1] + 0.5 * Rw * math.cos(theta_rad)],
        color='black', linewidth=line_width_wheel)
ax.text(wheels[0][0] + annotation_offset, wheels[0][1] - annotation_offset, "$m_4$", fontsize=20)
ax.plot([wheels[1][0] - 0.5 * Rw * -math.sin(theta_rad), wheels[1][0] + 0.5 * Rw * -math.sin(theta_rad)],
        [wheels[1][1] - 0.5 * Rw * math.cos(theta_rad), wheels[1][1] + 0.5 * Rw * math.cos(theta_rad)],
        color='black', linewidth=line_width_wheel)
ax.text(wheels[1][0] - annotation_offset * 3, wheels[1][1] - annotation_offset, "$m_3$", fontsize=20)
ax.plot([wheels[2][0] - 0.5 * Rw * -math.sin(phi_rad), wheels[2][0] + 0.5 * Rw * -math.sin(phi_rad)],
        [wheels[2][1] - 0.5 * Rw * -math.cos(phi_rad), wheels[2][1] + 0.5 * Rw * -math.cos(phi_rad)],
        color="black", linewidth=line_width_wheel)
ax.text(wheels[2][0] - annotation_offset * 3, wheels[2][1] + annotation_offset, "$m_2$", fontsize=20)
ax.plot([wheels[3][0] - 0.5 * Rw * math.sin(phi_rad), wheels[3][0] + 0.5 * Rw * math.sin(phi_rad)],
        [wheels[3][1] - 0.5 * Rw * -math.cos(phi_rad), wheels[3][1] + 0.5 * Rw * -math.cos(phi_rad)],
        color="black", linewidth=line_width_wheel)
ax.text(wheels[3][0] + annotation_offset, wheels[3][1] + annotation_offset, "$m_1$", fontsize=20)

# Print the vectors of the wheels
ax.arrow(*wheels[0], Rw * math.sin(theta_rad), Rw * math.cos(theta_rad), width=0.002, zorder=2, facecolor="gray")
ax.arrow(*wheels[1], Rw * math.sin(theta_rad), Rw * -math.cos(theta_rad), width=0.002, zorder=2, facecolor="gray")
ax.arrow(*wheels[2], Rw * -math.sin(phi_rad), Rw * -math.cos(phi_rad), width=0.002, zorder=2, facecolor="gray")
ax.arrow(*wheels[3], Rw * -math.sin(phi_rad), Rw * math.cos(phi_rad), width=0.002, zorder=2, facecolor="gray")

# Vertical lines going over the robot
ax.plot([-R, R], [0, 0], color="black", linestyle="--")

# Print the angles of the wheels
angle_m1 = Arc((0, 0), width=R, height=R, theta1=0, theta2=theta_deg, angle=-theta_deg)
angle_m2 = Arc((0, 0), width=R, height=R, theta1=0, theta2=theta_deg, angle=180)
angle_m3 = Arc((0, 0), width=1.1*R, height=1.1*R, theta1=0, theta2=phi_deg)
angle_m4 = Arc((0, 0), width=1.1*R, height=1.1*R, theta1=0, theta2=phi_deg, angle=180 - phi_deg)

# Print the symbol reflecting the amount of degrees
ax.text(wheels[0][0] * 1.2, wheels[0][1] / 3, "$\\theta = " + str(theta_deg) + "^{\\circ}$", fontsize=15)
ax.text(wheels[1][0] * 1.1, wheels[1][1] / 3, "$\\theta$", fontsize=15)
ax.text(wheels[2][0] / 1.4, wheels[2][1] / 3.3, "$\\phi$", fontsize=15)
ax.text(wheels[3][0] / 1.5, wheels[3][1] / 3.3, "$\\phi = " + str(phi_deg) + "^{\\circ}$", fontsize=15)

# Print the directions
ax.arrow(-1.2*R, -1.2*R, R/2, 0, width=0.001, facecolor=RTT_purple)
ax.text(-0.95*R, -(1.2-(1/16))*R, "$x$", fontsize=15)
ax.arrow(-1.2*R, -1.2*R, 0, R/2, width=0.001, facecolor=RTT_purple)
ax.text(-(1.2-(1/16))*R, -0.95*R, "$y$", fontsize=15)
ax.arrow(0, 0, 0, R/2, width=0.002, facecolor=RTT_purple)
ax.text(R/16, R/4, "$v$", fontsize=15)
ax.arrow(0, 0, R/2, 0, width=0.002, facecolor=RTT_purple)
ax.text(R/4, R/16, "$u$", fontsize=15)
omega = FancyArrowPatch((1.2 * R, 0), (0, 1.2 * R),
                        connectionstyle="arc3, rad="+str(6*R), facecolor=RTT_purple,
                        arrowstyle=f"Simple, head_width=10, head_length=20, tail_width=3")
ax.add_patch(omega)
ax.text(.95 * R, .95 * R, "$\omega$", fontsize=15)

ax.set_xlim(-R * 1.3, R * 1.3)
ax.set_ylim(-R * 1.3, R * 1.3)
ax.add_patch(arc)
ax.add_patch(angle_m1)
ax.add_patch(angle_m2)
ax.add_patch(angle_m3)
ax.add_patch(angle_m4)
ax.axis('off')
fig.tight_layout()
fig.show()
fig.savefig("topview.png")