print('Importing modules...', end='')
import tkinter as tk
import tinyik, math, time
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt
print('done')
class dotdict(dict):
    def __getattr__(self, attr):
        return self.get(attr, None)
    __setattr__= dict.__setitem__
    __delattr__= dict.__delitem__

class AutoEnum(object):
    def __init__(self):
        self.counter = 0
        self.d = {}
    def __getattr__(self, c):
        if c not in self.d:
            self.d[c] = self.counter
            self.counter += 1        
        return self.d[c]


class Controller(tk.Frame):
    def __init__(self, master, display):
        self.display = display
        self.history = []
        self.master = master
        self.master.minsize(300, 300)
        tk.Label(self.master, text='Joints').grid(row=0, column=0)
        self.joints_e = tk.Entry(self.master)
        self.joints_e.grid(row=0, column=1)
        tk.Label(self.master, text='Links').grid(row=1, column=0)
        self.links_e = tk.Entry(self.master)
        self.links_e.grid(row=1, column=1)
        tk.Label(self.master, text='Angles').grid(row=2, column=0)
        self.angles_e = tk.Entry(self.master)
        self.angles_e.grid(row=2, column=1)
        tk.Label(self.master, text='End effector').grid(row=3, column=0)
        self.ee_e = tk.Entry(self.master)
        self.ee_e.grid(row=3, column=1)
        tk.Button(self.master, text='QUIT', command=self.Q, bg='red').grid(row=4, column=0)
        self.update_b = tk.Button(self.master, text='Update', command=self.update)
        self.update_b.grid(row=4, column=1)
        tk.Button(self.master, text='History', command=self.history_w).grid(row=5, column=0)
        tk.Button(self.master, text='Animate', command=self.animate_w).grid(row=5, column=1)
    def Q(self):
        plt.close()
        self.master.destroy()
    def update(self):
        c1 = self.calc()
        c2 = self.display.calc(c1.angles, c1.joints, c1.links)
        print('c1out', c1)
        print('c2out', c2)
        self.history.append([c1, c2])
        self.display.draw(c2.x, c2.y, c2.z, c2.colors)
    def animate_w(self):
        win = tk.Toplevel()
        win.title('Animate')

        self.ainfo = dotdict()
        tk.Label(win, text='Frame IDs').grid(row=0, column=0)
        self.ainfo.ids = tk.Entry(win)
        self.ainfo.ids.grid(row=0, column=1)
        self.ainfo.animate = tk.Button(win, text='Animate', command=self.animate_frames)
        self.ainfo.animate.grid(row=1, column=1)
    def animate_frames(self):
        frame_ids = self.ainfo.ids.get().split(', ')
        for index, frame_id in enumerate(frame_ids):
            frame_ids[index] = int(frame_id)

        for index, frame_id in enumerate(frame_ids):
            print(frame_id)
            frame = self.history[frame_id][1]
            self.ainfo.animate.configure(text='Frame '+str(frame_id)+' Animation frame '+str(index)+'/'+str(len(frame_ids)))
            self.display.draw(frame.x, frame.y, frame.z, frame.colors)
            time.sleep(1)

        self.ainfo.animate.configure(text='Animate')
    def history_w(self):
        win = tk.Toplevel()
        win.title('History')
        
        self.hinfo = dotdict()
        tk.Label(win, text='Frame ID').grid(row=0, column=0)
        self.hinfo.id = tk.Entry(win)
        self.hinfo.id.grid(row=0, column=1)
        tk.Label(win, text='Joints').grid(row=1, column=0)
        self.hinfo.joints = tk.Entry(win)
        self.hinfo.joints.grid(row=1, column=1)
        tk.Label(win, text='Links').grid(row=2, column=0)
        self.hinfo.links = tk.Entry(win)
        self.hinfo.links.grid(row=2, column=1)
        tk.Label(win, text='Angles').grid(row=3, column=0)
        self.hinfo.angles = tk.Entry(win)
        self.hinfo.angles.grid(row=3, column=1)
        tk.Label(win, text='End effector').grid(row=4, column=0)
        self.hinfo.ee = tk.Entry(win)
        self.hinfo.ee.grid(row=4, column=1)
        tk.Button(win, text='Load', command=self.load_h).grid(row=5, column=0)
        tk.Button(win, text='Draw', command=self.draw_h).grid(row=5, column=1)
    def load_h(self):
        instance = self.history[int(self.hinfo.id.get())][0]
        
        self.hinfo.joints.delete(0, 50)
        self.hinfo.joints.insert(0, instance.str_joints)
        self.hinfo.links.delete(0, 50)
        self.hinfo.links.insert(0, instance.str_links)
        self.hinfo.angles.delete(0, 50)
        self.hinfo.angles.insert(0, instance.str_angles)
        self.hinfo.ee.delete(0, 50)
        self.hinfo.ee.insert(0, instance.str_ee)
    def draw_h(self):
        instance = self.history[int(self.hinfo.id.get())][1]
        self.display.draw(instance.x, instance.y, instance.z, instance.colors)
    def calc(self):
        joints = self.joints_e.get().split(', ')
        links = self.links_e.get().split(', ')
        angles = self.angles_e.get().split(', ')
        ee = self.ee_e.get()
        argument = []
        for i, joint in enumerate(joints):
            argument.append(joint)
            dem = links[i][-1]
            num = float(links[i][:-1])
            if dem == 'x':
                links[i] = [num, 0, 0]
            elif dem == 'y':
                links[i] = [0, num, 0]
            else:
                links[i] = [0, 0, num]
            argument.append(links[i])
        arm = tinyik.Actuator(argument)
        if angles == ['']:
            angles = arm.angles
        else:
            for i in range(0, len(angles)):
                angles[i] = np.deg2rad(float(angles[i]))
            arm.angles = angles
        if ee == '':
            ee = arm.ee
        else:
            ee = ee.split(', ')
            ee = np.array([float(ee[0]), float(ee[1]), float(ee[2])])
            arm.ee = ee
        angles = np.round(np.rad2deg(arm.angles))
        ee = np.round(arm.ee)
        str_angles = str(angles.tolist()).lstrip('[').rstrip(']')
        self.angles_e.delete(0, 50)
        self.angles_e.insert(0, str_angles)
        str_ee = str(ee.tolist()).lstrip('[').rstrip(']')
        self.ee_e.delete(0, 50)
        self.ee_e.insert(0, str_ee)
        return dotdict({'arm': arm, 'joints': joints, 'links': links, 'ee': ee, 'angles': angles,
                        'str_links': self.links_e.get(), 'str_joints': self.joints_e.get(),
                        'str_angles': str_angles, 'str_ee': str_ee})

class Display:
    def __init__(self):
        self.fig = plt.figure()
        self.ax = self.fig.gca(projection='3d')
        plt.show(block=False)
    def calc(self, angles, joints, links):
        print('c2in', angles, joints, links)
        x_list = [0]
        y_list = [0]
        z_list = [0]
        colors = []
        for i in range(0, len(angles)):
            x = 0
            y = 0
            z = 0
            theata = sum(angles[:i+1])
            cos = math.cos(math.radians(theata))
            sin = math.sin(math.radians(theata))
            joint = joints[i]
            link = links[i]
            if link[0] != 0:
                dem = 'x'
                num = link[0]
            elif link[1] != 0:
                dem = 'y'
                num = link[1]
            else:
                dem = 'z'
                num = link[2]

            if joint == 'x':
                if dem == 'y':
                    y = cos*num
                    z = sin*num
                else:
                    z = cos*num
                    y = sin*num
            elif joint == 'y':
                if dem == 'x':
                    x = cos*num
                    z = sin*num
                else:
                    z = cos*num
                    x = sin*num
            else:
                if dem == 'x':
                    x = cos*num
                    y = sin*num
                else:
                    y = cos*num
                    x = sin*num
            x_list.append(x_list[-1] + x)
            y_list.append(y_list[-1] + y)
            z_list.append(z_list[-1] + z)
            colors.append({'x': 'red',
                           'y': 'green',
                           'z': 'blue'}[joint])
        print('internal c2out', x_list, y_list, z_list, colors)
        return dotdict({'x': x_list, 'y': y_list, 'z': z_list, 'colors': colors})
    def draw(self, x, y, z, c='red'):
        self.ax.cla()
        print(x, y, z, c)
        self.ax.plot(x, y, z)
        c.append('orange')
        self.ax.scatter(x, y, z, c=c, depthshade=False)
        plt.show(block=False)

if __name__ == '__main__':
    display = Display()
    root = tk.Tk()
    win = Controller(root, display)
    win.master.mainloop()