from PyQt4 import QtGui, QtCore, uic
import sip
import maya.OpenMayaUI as mui
import functools

ptr = mui.MQtUtil.mainWindow()
maya_window = sip.wrapinstance(long(ptr), QtCore.QObject)
maya_timeline = maya_window.findChild(QtGui.QWidget, "timeControl1")

class TimeLineMarkers(QtGui.QWidget):
    def __init__(self, parent = maya_timeline):
        super(TimeLineMarkers, self).__init__(parent)
        self.layout = QtGui.QVBoxLayout()
        self.setLayout(self.layout)
        self.layout.setContentsMargins(0,0,0,0)
        self.marks = []
    def paintEvent(self, event):
        if not self.marks: return
        color = QtGui.QColor(255, 0, 0, 0)
        painter = QtGui.QPainter(self)
        pen = QtGui.QPen(QtGui.QColor(200,200,255,100))
        pen.setWidth(3)
        painter.setPen(pen)
        for m in self.marks:
            line = QtCore.QLineF(QtCore.QPointF(m,0), QtCore.QPointF(m, 100))
            painter.drawLine(line)

    def context_menu(self, point):
        menu = QtGui.QMenu(self)
        menu.addAction("add mark", functools.partial(self.cm_add_mark, point))
        menu.addAction("clear all", self.cm_clear)
        menu.popup(self.mapToGlobal(point))

    def cm_add_mark(self, *args):
        pos = args[0]
        self.marks.append(pos.x())
        self.update()

    def cm_clear(self):
        self.marks = []
        self.update()

class ShiftKeyFilter(QtCore.QObject):
    shift_down = False
    def eventFilter(self, obj, event):
        if event.type() == QtCore.QEvent.KeyPress:
            if event.key() == 16777248:
                self.shift_down = True
            return False
        if event.type() == QtCore.QEvent.KeyRelease:
            if event.key() == 16777248:
                self.shift_down = False
            return False
        return False

class RightClickFilter(QtCore.QObject):
    def __init__(self, shift_key_filter, context_func):
        super(RightClickFilter, self).__init__()
        self.shift_key_filter = shift_key_filter
        self.context_func = context_func
    def eventFilter(self, obj, event):
        if event.type() == QtCore.QEvent.MouseButtonPress:
            if self.shift_key_filter.shift_down:
                self.shift_key_filter.shift_down = False
                self.context_func(event.pos())
            return False
        return False

def main():
    global time_line_markers
    try:
        time_line_markers.hide();time_line_markers.setParent(None)
    except:
        pass
    timeline_layout = None
    timeline_layout = maya_timeline.layout()
    if timeline_layout == None:
        timeline_layout = QtGui.QVBoxLayout(maya_timeline)
        timeline_layout.setContentsMargins(0,0,0,0)
        maya_timeline.setLayout(timeline_layout)

    timeline_markers = TimeLineMarkers()
    timeline_layout.addWidget(timeline_markers)

    #install event filters to get a shift+click menu on the time line
    global shift_key_filter
    global right_click_filter
    shift_key_filter = ShiftKeyFilter()
    right_click_filter = RightClickFilter(shift_key_filter, timeline_markers.context_menu)
    try:
        maya_window.removeEventFilter(shift_key_filter)
        maya_timeline.removeEventFilter(right_click_filter)
    except:
        pass
    maya_window.installEventFilter(shift_key_filter)
    maya_timeline.installEventFilter(right_click_filter)

main()