import matplotlib.pyplot as plt
import numpy as np
class ArmVisualizer2Dsom(object):
    def __init__(self, arm, problem):
        self.arm = arm
        self.problem = problem
        

        fig=plt.figure()
        cid = fig.canvas.mpl_connect('button_press_event', self.onclick)
        cid_k=fig.canvas.mpl_connect('key_press_event', self.onKeyPress)

        self.subplot_activation={}
        self.subplot_activation[arm.l]=plt.subplot(241)
        arm.l.plot_activation('green')
        self.subplot_activation[arm.s0]=plt.subplot(242)
        arm.s0.plot_activation('green')

        self.subplot_activation[arm.s1]=plt.subplot(243)
        arm.s1.plot_activation('green')

#        self.subplot_activation[arm.c01]=plt.subplot(244)#,projection='3d')

        self.subplot_arm_schematic=plt.subplot(245)
        self.subplot_arm_schematic.set_xlim(self.arm.l.limits[0,0],self.arm.l.limits[0,1])
        self.subplot_arm_schematic.set_ylim(self.arm.l.limits[1,0],self.arm.l.limits[1,1])
        self.subplot_arm_target_line=plt.Line2D([0.0,0.0],[0.0,0.0])
        self.subplot_arm_schematic.add_line(self.subplot_arm_target_line)

        
        

        #self.subplot_connections_0=plt.subplot(246)
        #self.subplot_connections_0.my_colorbar=None
        #self.subplot_connections_1=plt.subplot(247)
        #self.subplot_connections_1.my_colorbar=None


        plt.show()






    def draw_pose_2d(self,pose, color, axes):

        seg_ends=self.arm.segment_ends_from_pose((pose[0],pose[1]))#-pose[0]))
        axes.plot([0.0,seg_ends[0,0],seg_ends[2,0]],[0.0,seg_ends[0,1],seg_ends[2,1]])

    def draw_target_line(self,color,field, axes=None):
        if not axes:
            self.subplot_activation[field].plot([0.0,self.arm.target_position[field][0]],[0.0,self.arm.target_position[field][1]])
        else:
            axes.plot([0.0,self.arm.target_position[field][0]],[0.0,self.arm.target_position[field][1]])
#        def draw_connections_s0_l(self, axes, idx_of_nearest_neuron):
#            W_nearest_l=self.arm.W[idx_of_nearest_neuron,self.arm.start_idx[arm.l]:self.arm.end_idx[arm.l]]
#            W_nearest_l=np.reshape(W_nearest_l,(self.arm.l_y_number_of_neurons,self.arm.l_x_number_of_neurons))
#            W_nearest_l=np.flipud(W_nearest_l)
#            return axes.matshow(W_nearest_l,cmap=plt.cm.gray)

    def draw_connections(self, axes, idx_of_nearest_neuron, f0, f1):
        W_nearest_l=self, arm.W[self.arm.start_idx[f0]:self.arm.end_idx[f0],self.arm.start_idx[f1]+idx_of_nearest_neuron]
        xs,ys=zip(*f0.places)
        axes.cla()
        scatter_plot=axes.scatter(xs,ys,s=50,c=W_nearest_l,cmap=plt.cm.get_cmap('RdYlBu'))
        if not axes.my_colorbar:
            axes.my_colorbar=plt.colorbar(scatter_plot,ax=axes)
        axes.my_colorbar.set_clim(vmin=min(W_nearest_l),vmax=max(W_nearest_l))
        axes.my_colorbar.draw_all()     

    def onclick(self, event):
        print 'button=%d, x=%d, y=%d, xdata=%f, ydata=%f'%(
            event.button, event.x, event.y, event.xdata, event.ydata)

        click_position=np.array([event.xdata,event.ydata])
        if self.problem=="kinematics_inverse":
            if event.inaxes == self.subplot_activation[self.arm.l]:
                self.arm.target_position[self.arm.l]=click_position
                self.draw_target_line('red',self.arm.l)
                


                #plot target line
#                self.subplot_arm_schematic.cla()
#                self.subplot_arm_schematic.set_xlim(self.arm.l.limits[0,0],self.arm.l.limits[0,1])
#                self.subplot_arm_schematic.set_ylim(self.arm.l.limits[1,0],self.arm.l.limits[1,1])
                

        elif self.problem=="kinematics_forward":
            if event.inaxes == self.subplot_activation[self.arm.s0]:
                self.arm.target_position[self.arm.s0]=click_position
                self.draw_target_line('red',self.arm.s0)
            if event.inaxes == self.subplot_activation[self.arm.s1]:
                self.arm.target_position[self.arm.s1]=click_position
                self.draw_target_line('red',self.arm.s1)
        #if event.inaxes == self.subplot_activation[arm.s0]:
            #find next neuron
         #   idx_of_nearest_neuron=arm.s0.idx_of_nearest_neuron(click_position)
          #  self.draw_connections(self.subplot_connections_0,idx_of_nearest_neuron,arm.l,arm.s0)
            #plt.colorbar(connection_plot)

        #if event.inaxes == subplot_pf_s1:
            #find next neuron
         #   idx_of_nearest_neuron=arm.s1.idx_of_nearest_neuron(click_position)
          #  draw_connections(subplot_connections_1,idx_of_nearest_neuron,arm.l,arm.s1)

        plt.draw()


    def onKeyPress(self, event):
        if event.key == 'up':
            if self.problem=="kinematics_inverse":
                self.arm.activate_by_target_position(self.arm.l,method="exponential")
            if self.problem=="kinematics_forward":
                self.arm.activate_by_target_position(self.arm.s0,method="exponential")
                self.arm.activate_by_target_position(self.arm.s1,method="exponential")
            
            self.arm.simulate_one_step()
        if event.key == 'left':
            for i in range(10):
                if self.problem=="kinematics_inverse":
                    self.arm.activate_by_target_position(self.arm.l,method="exponential")
                if self.problem=="kinematics_forward":
                    self.arm.activate_by_target_position(self.arm.s0,method="exponential")
                    self.arm.activate_by_target_position(self.arm.s1,method="exponential")


                self.arm.simulate_one_step()


        if event.key == 'right':
            self.arm.simulate_one_step()

        if event.key == 'r':
            #randomize activations:
            self.arm.randomize_activations(1.0)

        if event.key == 'z':
            #randomize activations:
            self.arm.activations[:]=0.0

        if event.key == '1':
            self.arm.activations[:]=0.0
            self.arm.s0.activations[0]=1.0
            self.arm.s1.activations[3]=1.0

        if event.key == 'c':
            self.arm.activate_by_som_and_target()

        for field, subplot in self.subplot_activation.items():
            subplot.cla()
            field.plot_activation('green',subplot)

        self.subplot_arm_schematic.cla()
        self.subplot_arm_schematic.set_xlim(self.arm.l.limits[0,0],self.arm.l.limits[0,1])
        self.subplot_arm_schematic.set_ylim(self.arm.l.limits[1,0],self.arm.l.limits[1,1])
        most_probable_pose=self.arm.get_most_probable_pose()

        

        if self.problem=="kinematics_inverse":
            if self.arm.l in self.arm.target_position.keys():
                self.draw_target_line('red',self.arm.l)
                self.draw_target_line('red',self.arm.l, self.subplot_arm_schematic)
        if self.problem=="kinematics_forward":
            if self.arm.s0 in self.arm.target_position.keys():
                self.draw_target_line('red',self.arm.s0)
            if self.arm.s1 in self.arm.target_position.keys():
                self.draw_target_line('red',self.arm.s1)
        self.draw_pose_2d(most_probable_pose,'green', self.subplot_arm_schematic)
        plt.draw()

        if event.key == "w":
            plt.figure(2)
            W_ax=plt.matshow(self.arm.W,cmap=plt.cm.gray,fignum=2)
            plt.colorbar(W_ax)
            plt.show()
            #plt.figure(1)



