# -*- coding: utf-8 -*-
"""
A simple menu based on matplotlib figure.legend
"""
import matplotlib.pyplot as plt
import matplotlib.lines as mline
from matplotlib.text import Text
from math import ceil
from PyQt5 import QtWidgets   # , QtGui, QtCore

class Menu:
    '''a simple matplotlib menu based on Legend
    
    Steps:
        (1) menu = Menu()
        (2) use menu as a decorator to add menu items
        (3) call create_menu(fig, ncol) to actually create the menu
        (4) menu.add_qt_toolbar_button()
    '''
    def __init__(self):
        self.label_funcs = {}  # 
        
    def __call__(self, label):
        'may need to refine to distinguish function and method'
        def wrap(f):
            def wrapped_f(*args):
                f(*args)

            self.add_menu_item(label, wrapped_f)
            return wrapped_f
        return wrap    
    
    def add_menu_item(self, label, func):
        'just collect, because Legend does not allow incrementally add entries'
        self.label_funcs[label] = func
        
    # the following functions needs figure
    def create_menu(self, fig, ncol):
        'create menu'
        self.fig = fig
        dummy = mline.Line2D((0, 1), (0, 1) )
        self.num = num = len(self.label_funcs)
        self.ncol = ncol
        self.nrow = nrow = ceil(num/ncol)

        self.leg = leg = fig.legend([dummy]*num, 
                     self.label_funcs.keys(),
             loc='upper right',
             ncol=ncol,  
             borderaxespad=0., 
             handlelength=0, handleheight=0,
             fancybox=True, shadow=True,
             framealpha=0.3,
             # title='Menu',
             )
        leg.draggable(True)  
        leg.set_visible(True)

        for t in leg.get_texts():
            t.set_picker(5)  # 5 pts tolerance

        # inner function to create_menu
        last_pick = None
        def onpick(event):
            nonlocal last_pick
            if not isinstance(event.artist, Text):
                return
           
            # restore lp
            if last_pick:
                last_pick.set_bbox(None)

            # set new pick
            last_pick = legtext = event.artist
            label = legtext.get_text()
            legtext.set_backgroundcolor('#87CEEBB0')
            self.label_funcs[label]()
            fig.canvas.draw_idle()
        
        self.cid_pick = fig.canvas.mpl_connect('pick_event', onpick)
        
        last_hover = None
        def onhover(event):
            nonlocal last_hover
            if not self.leg.get_visible(): return
            over, junk = self.leg.contains(event)
            if not over: return
            
            x0, y0, width, height = self.leg.get_window_extent().bounds
            idx = int((event.x - x0)/width * self.ncol)
            idy = int(self.nrow - (event.y - y0)/height * self.nrow)
            id = idx * self.nrow + idy
            
            t = self.leg.get_texts()[id]
            
            if t == last_hover:
                return
            else:
                if last_hover:
                    last_hover.set_alpha(1.0)
                last_hover = t
                t.set_alpha(0.3)
                    
            fig.canvas.draw_idle()
         
        # matplotlib 2.1.0 busy cursor makes hovering annoying
        self.cid_hover = fig.canvas.mpl_connect('motion_notify_event', onhover)

    def add_qt_toolbar_button(self):             
        # add a toolbar button to toggle the menu
        # inner function for toolbutton
        def onoff(label):
            if self.leg.get_visible():
                self.leg.set_visible(False)
            else:
                self.leg.set_visible(True)
            fig.canvas.draw_idle()
            
        toolbar = fig.canvas.toolbar
        mtb = QtWidgets.QToolButton(toolbar)
        style = QtWidgets.QStyle
        stdicon = mtb.style().standardIcon
        mtb.setIcon(stdicon(style.SP_VistaShield))
        #mtb.setIconSize(QtCore.QSize(48, 48))
        mtb.setText('Menu')
        mtb.setToolTip('Menu')
        toolbar.addWidget(mtb)
        mtb.clicked.connect(onoff)
        
if __name__ == '__main__':
    import numpy as np
    
    menu = Menu()
    
    fig, ax = plt.subplots()
    plt.subplots_adjust(bottom=0.2)
    t = np.arange(0.0, 1.0, 0.001)
    s = np.sin(2*np.pi*5*t)
    line, = plt.plot(t, s, lw=2)

    
    @menu('Red')
    def abc():
        line.set_color('red')
    
    @menu('Blue')
    def cde():
        line.set_color('blue')

    @menu('Black')
    def fgh():
        line.set_color('black')

    @menu('Thick')
    def f1():
        line.set_linewidth(5)

    @menu('3 pts')
    def f2():
        line.set_linewidth(3)

    @menu('Thin')
    def f3():
        line.set_linewidth(1)
        
    @menu('Dash')
    def fd():
        line.set_linestyle('--')

    @menu('Solid')
    def fs():
        line.set_linestyle('-')     
        
    @menu('Dotted')
    def fo():
        line.set_linestyle(':')  
        
    menu.create_menu(fig, ncol=3)
    menu.add_qt_toolbar_button()
       
    plt.show()