from typing import Tuple

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backend_bases import MouseButton

def circle(rad, px, py, linestyle, ax):
    t = np.linspace(0, 2*np.pi, 10)
    ax.plot(rad * np.cos(t) + px, rad * np.sin(t) + py, color='k', linestyle=linestyle)

def plotSeg(f, t, linestyle, ax):
    d = t - f
    l = np.linalg.norm(d)
    r1 = l / 10
    r2 = l / 15
    r3 = l / 7
    d = d / l

    # Plot the motor joint
    circle(l / 10, np.real(f), np.imag(f), linestyle, ax)
    circle(l / 15, np.real(t), np.imag(t), linestyle, ax)

    # Plot the link
    pts = np.array([f + r1*1j*d, t + r2*1j*d])
    ax.plot(np.real(pts), np.imag(pts), color='k', linestyle=linestyle)
    pts = np.array([f - r1*1j*d, t - r2*1j*d])
    ax.plot(np.real(pts), np.imag(pts), color='k', linestyle=linestyle)

    # Plot the joint axis
    pts = np.array([t, t + r3 * d])
    ax.plot(np.real(pts), np.imag(pts), color='r', linestyle=linestyle)
    pts = np.array([t, t + 1j * r3 * d])
    ax.plot(np.real(pts), np.imag(pts), color='g', linestyle=linestyle)

class Robot2D:
    def __init__(
        self,
        figsize: Tuple[int, int] = (5, 5),
        ):

        # Initialize figure once
        self.fig = plt.figure(figsize=figsize)
        self.ax = self.fig.add_subplot()
        self.fig.canvas.mpl_connect('close_event', self._on_close)
        self.terminate = False
    
    def _on_close(self, event):
        """Event handler, should the figure be manually closed.
        """
        self.terminate = True

    def draw_point(self, pos: np.ndarray):
        """Draw goal position (from ginput).
        """
        if not self.terminate:
            self.ax.scatter(pos[0], pos[1], marker="*")      # Draw the point

    def plotRobot2D( self, angles: np.ndarray, link_lengths: np.ndarray, linestyle: str = "-",):
        """Plot the robot.
        """
        # Restrict workspace & figure size
        lim = np.sum(link_lengths) + 0.1      # Small overflow
        self.ax.set_xlim([-lim, lim])
        self.ax.set_ylim([-lim, lim])
        self.ax.set_adjustable("datalim")

        f = 0
        t = 0
        for i in range(len(link_lengths)):
            # Forward kinematics
            t = t + angles[i]
            p = f + (np.cos(t) + 1j*np.sin(t)) * link_lengths[i]
            plotSeg(f, p, linestyle, self.ax)
            f = p

        # Non-blocking show
        plt.show(block=False)