#!/usr/bin/env python

import wx

import numpy as np


data = [
    (-1, 0, 0, 350),  # root
    (0, 60, 1, 220),  # 1
    (0, 60, 2, 140),
    (0, 60, 3, 180),
    (0, 100, 4, 140),
    (0, 100, 5, 140),  # 5
    (0, 100, 6, 320),
    (0, 200, 7, 240),
    (1, 180, 8, 260),
    (3, 140, 9, 240),
    (0, 200, 10, 260), # 10
    (1, 200, 11, 240),
    (11, 220, 12, 300),
    (11, 220, 13, 320),
    (3, 100, 14, 260),
]


def get_order(source, parent_level):
    """Find restarts with given parent level.

    Building the list from bottom right up, so first restart at the largest
    frame number is first, then later restarts at the same frame number after
    that. Checkpoints that occur at earlier frames then are added, proceeding
    down to frame zero.
    """
    subset_lookup = {}
    for entry in [d for d in source if d[0] == parent_level]:
        start = entry[1]
        if start not in subset_lookup:
            subset_lookup[start] = []
        subset_lookup[start].append(entry)
    processing_order = []
    for items in reversed(list(subset_lookup.values())):
        for item in items:
            processing_order.append(item)
            processing_order.extend(get_order(source, item[2]))
    return processing_order


class RestartLine:
    def __init__(self, data, level):
        self.start_frame = data[1]
        self.end_frame = data[3]
        self.parent = data[0]
        self.restart_number = data[2]
        try:
            current_largest_level = max(level[self.start_frame:self.end_frame + 1])
        except ValueError:
            current_largest_level = 0
        self.level = current_largest_level + 1
        level[self.start_frame:self.end_frame + 1] = self.level
        self.interesting_frames = {
            self.start_frame: "Checkpoint from run {self.parent}",
            self.end_frame: "Last frame",
        }

    def __repr__(self):
        return f"{self.restart_number}: {self.start_frame}->{self.end_frame} @ {self.level}"

    def is_frame_in_line(self, frame):
        return frame >= self.start_frame and frame <= self.end_frame


class RestartLines:
    def __init__(self, processing_order):
        self.lines = [None] * len(processing_order)
        self.last_frame = max([d[3] for d in processing_order])
        self.highest_level = 0
        self.generate_lines(processing_order)

    def __str__(self):
        txt_lines = []
        for line in self.lines:
            txt_lines.append(repr(line))
        return "\n".join(txt_lines)

    def __iter__(self):
        for line in self.lines:
            yield line

    def __getitem__(self, index):
        return self.lines[index]

    def generate_lines(self, processing_order):
        level = np.zeros(self.last_frame, dtype=np.int16) - 1
        for d in processing_order:
            line = RestartLine(d, level)
            self.lines[line.restart_number] = line
            if line.level > self.highest_level:
                self.highest_level = line.level
            #print(level)
            # print("generate", line, line.level)

    def find_line(self, frame_number, level):
        for line in self.lines:
            if frame_number >= line.start_frame and frame_number <= line.end_frame and level == line.level:
                return line
        return None

    def find_parent(self, child_line):
        return self.lines[child_line.parent]


wxEVT_CHECKPOINT_RESTORE = wx.NewEventType()
EVT_CHECKPOINT_RESTORE = wx.PyEventBinder(wxEVT_CHECKPOINT_RESTORE, 1)
wxEVT_CHECKPOINT_HOVER = wx.NewEventType()
EVT_CHECKPOINT_HOVER = wx.PyEventBinder(wxEVT_CHECKPOINT_HOVER, 1)
wxEVT_CHECKPOINT_SELECTED = wx.NewEventType()
EVT_CHECKPOINT_SELECTED = wx.PyEventBinder(wxEVT_CHECKPOINT_SELECTED, 1)
wxEVT_CHECKPOINT_RANGE_SELECTED = wx.NewEventType()
EVT_CHECKPOINT_RANGE_SELECTED = wx.PyEventBinder(wxEVT_CHECKPOINT_RANGE_SELECTED, 1)


class CheckpointTreeEvent(wx.PyCommandEvent):
    """This event class is used for all events generated by the CheckpointTree
    """
    def __init__(self, type, obj, mouse_info):
        wx.PyCommandEvent.__init__(self, type)
        self.SetEventObject(obj)
        self.SetId(obj.GetId())
        self.mouse_info = mouse_info

    def GetLine(self):
        """The RestartLine of the event
        """
        return self.mouse_info.nearest_line

    def GetRestartNumber(self):
        """The restart number of the event; only valid if there is a line
        associated with this event, otherwise will raise AttributeError
        """
        return self.mouse_info.restart_number

    def GetFrameNumber(self):
        """The frame number of the event
        """
        return self.mouse_info.frame_number


class CheckpointTreeRangeEvent(wx.PyCommandEvent):
    """This event class is used for all events generated by the CheckpointTree
    """
    def __init__(self, type, obj, start, end):
        wx.PyCommandEvent.__init__(self, type)
        self.SetEventObject(obj)
        self.SetId(obj.GetId())
        self.start = start
        self.end = end


class CheckpointTreeMouseInfo:
    def __init__(self, tree, evt):
        self.tree = tree
        self.nearest_line = None
        self.calc_mouse_info(evt)

    @property
    def is_near(self):
        return self.nearest_line is not None

    @property
    def restart_number(self):
        return self.nearest_line.restart_number

    def calc_mouse_info(self, evt):
        tree = self.tree
        ex, ey = evt.GetX(), evt.GetY()
        sx, sy = tree.GetViewStart()
        x, y = tree.CalcUnscrolledPosition(ex, ey)
        self.x, self.y = x, y
        size = tree.GetVirtualSize()
        # print(ex, sx, x)

        self.frame_number = (x - tree.width_border + tree.x_scale // 2) // tree.x_scale
        self.level = ((tree.virtual_height - y) - tree.height_border + tree.level_height // 2) // tree.level_height
        x1 = tree.frame_to_x(self.frame_number)
        y1 = tree.level_to_y(self.level)
        self.x_hit = abs(x1 - x)
        self.y_hit = abs(y1 - y)
        if self.y_hit < tree.y_hit:
            # print(frame_number, level)
            self.nearest_line = tree.restart_lines.find_line(self.frame_number, self.level)
        else:
            self.nearest_line = None


class CheckpointTree(wx.ScrolledWindow):
    def __init__(self, parent, source):
        wx.ScrolledWindow.__init__(self, parent, -1)

        self.width_border = 10
        self.height_border = 10
        self.x_scale = 1
        self.level_height = 10
        self.line_width = 3
        self.marker_size = 3
        self.x_hit = 5
        self.y_hit = 5
        self.over_line = None
        self.highlight_over_line = False
        self.scroll_chunk = 20
        self.mouse_mode = None
        self.initial_mouse_state = None
        self.current_mouse_state = None
        self.virtual_height = -1
        self.virtual_width = -1
        self.compute_lines(source)

        self.bg_color = wx.WHITE
        self.line_color = wx.BLUE
        self.marker_color = wx.Colour(255, 255, 255)
        self.cursor_color = wx.Colour(255, 0, 255)
        self.create_dc_stuff()

        self.SetBackgroundColour(self.bg_color)

        self.SetScrollRate(self.scroll_chunk, self.scroll_chunk)

        self.Bind(wx.EVT_PAINT, self.on_paint)
        self.Bind(wx.EVT_SIZE, self.on_size)
        self.Bind(wx.EVT_LEFT_DOWN, self.on_left_down)
        self.Bind(wx.EVT_MOTION, self.on_motion)
        self.Bind(wx.EVT_LEFT_UP, self.on_left_up)

    def compute_lines(self, source=None):
        if source is None:
            source = self.source
        self.source = source
        data = source.get_restart_summary()
        # print("DATA", data)
        parent_level = 0
        data.sort()
        order = get_order(data, parent_level)
        order[0:0] = [data[0]]
        # from pprint import pprint
        # pprint(order)
        lines = RestartLines(order)
        # print("restart_lines", lines)
        self.set_lines(lines)
        i = self.source.current_frame_number
        if i is not None:
            # print("current frame number:", i)
            dx = self.GetScrollPos(wx.HORIZONTAL)
            if i > dx:
                self.Scroll(i, -1)

    def recalc_view(self):
        self.compute_lines()

    def refresh_view(self):
        self.Refresh()

    def set_lines(self, lines):
        self.restart_lines = lines
        self.max_width = lines.last_frame * self.x_scale + 2 * self.width_border
        self.max_height = self.level_height * lines.highest_level + 2 * self.height_border
        if self.max_width > self.virtual_width:
            self.virtual_width = self.max_width
        if self.max_height > self.virtual_height:
            self.virtual_height = self.max_height
        self.SetVirtualSize((self.max_width, self.max_height))

    def create_dc_stuff(self):
        self.pen = wx.Pen(self.line_color, self.line_width)
        self.hover_pen = wx.Pen(self.line_color, self.line_width * 2 + 1)
        self.marker_brush = wx.Brush(self.marker_color)
        self.marker_pen = wx.Pen(self.line_color, 1)
        self.cursor_brush = wx.Brush(self.cursor_color)

    def on_size(self, event):
        size = self.GetVirtualSize()
        self.virtual_width = size.x
        self.virtual_height = size.y

    def on_paint(self, event):
        dc = wx.PaintDC(self)
        x0, y0 = self.CalcUnscrolledPosition (0, 0)
        dc.SetLogicalOrigin(x0, y0)

        dc.SetPen(self.pen)
        for line in self.restart_lines:
            x1 = self.frame_to_x(line.start_frame)
            x2 = self.frame_to_x(line.end_frame)
            y = self.level_to_y(line.level)
            if self.highlight_over_line and line == self.over_line:
                dc.SetPen(self.hover_pen)
                dc.DrawLine(x1, y, x2, y)
            else:
                dc.SetPen(self.pen)
                dc.DrawLine(x1, y, x2, y)
            dc.SetPen(self.marker_pen)
            dc.SetBrush(self.marker_brush)
            for interesting_frame in line.interesting_frames.keys():
                xi = self.frame_to_x(interesting_frame)
                dc.DrawCircle(xi, y, self.marker_size)

            if line.restart_number == 0:
                continue

            # draw vertical line connecting to parent restart
            parent_line = self.restart_lines[line.parent]
            parent_y = self.level_to_y(parent_line.level)
            dc.SetPen(self.pen)
            dc.DrawLine(x1, y, x1, parent_y)

        # highlighting of selected lines
        if self.current_mouse_state is not None:
            dc.SetPen(self.hover_pen)
            start, end = self.initial_mouse_state, self.current_mouse_state
            if start.frame_number > end.frame_number:
                start, end = end, start
            current_frame = end.frame_number
            leftmost_frame = start.frame_number
            line = end.nearest_line
            while current_frame > leftmost_frame:
                if line.is_frame_in_line(leftmost_frame):
                    left_frame = leftmost_frame
                else:
                    left_frame = line.start_frame
                x1 = self.frame_to_x(left_frame)
                x2 = self.frame_to_x(current_frame)
                y = self.level_to_y(line.level)
                dc.DrawLine(x1, y, x2, y)
                current_frame = left_frame
                if current_frame > leftmost_frame:
                    line = self.restart_lines.find_parent(line)

        # current frame & restart number
        i = self.source.current_frame_number
        if i is not None:
            xc = self.frame_to_x(i)
            restart = self.source.get_restart_of_frame(i)
            line = self.restart_lines[restart.restart_number]
            yc = self.level_to_y(line.level)
            dc.SetPen(self.marker_pen)
            dc.SetBrush(self.cursor_brush)
            dc.DrawCircle(xc, yc, self.marker_size)

    def frame_to_x(self, frame_number):
        return self.width_border + (self.x_scale * frame_number)

    def level_to_y(self, level):
        return self.virtual_height - (self.height_border + level * self.level_height)

    def on_left_down(self, evt):
        mouse = CheckpointTreeMouseInfo(self, evt)
        if mouse.is_near:
            evt = CheckpointTreeEvent(wxEVT_CHECKPOINT_SELECTED, self, mouse)
            self.GetEventHandler().ProcessEvent(evt)
            self.initial_mouse_state = mouse
            self.current_mouse_state = None
            self.mouse_mode = "range"

    def on_motion(self, evt):
        mouse = CheckpointTreeMouseInfo(self, evt)
        if mouse.is_near:
            if self.mouse_mode == "range":
                self.current_mouse_state = mouse
                evt = CheckpointTreeRangeEvent(wxEVT_CHECKPOINT_RANGE_SELECTED, self, self.initial_mouse_state, mouse)
                self.GetEventHandler().ProcessEvent(evt)
            else:
                evt = CheckpointTreeEvent(wxEVT_CHECKPOINT_HOVER, self, mouse)
                self.GetEventHandler().ProcessEvent(evt)
                if self.over_line != mouse.nearest_line:
                    # print("REFRESHING")
                    self.over_line = mouse.nearest_line
                    wx.CallAfter(self.Refresh)

    def on_left_up(self, evt):
        self.mouse_mode = None
        evt.Skip()


#For testing 
if __name__ == '__main__':
    app = wx.App()
    frame = wx.Frame(None, -1, "Test CheckpointTree", size=(500,400))

    from pprint import pprint

    data.sort()
    pprint(data)

    min_frame = min([d[1] for d in data])
    pprint(min_frame)

    max_frame = max([d[3] for d in data])
    pprint(max_frame)

    parent_level = 0
    order = get_order(data, parent_level)
    order[0:0] = [data[0]]
    pprint(order)

    lines = RestartLines(order)
    print("restart_lines", lines)

    class TestSource:
        current_frame_number = None
        def get_restart_summary(self):
            return list(data)

    scroll = CheckpointTree(frame, TestSource())
    frame.Show() 
    app.MainLoop()
