#
# dxplore.py - GUI for interactive exploration of digital waveforms
#
# Copyright (C) 2008, 2009 by OpenMoko, Inc.
# Written by Werner Almesberger <werner@openmoko.org>
# All Rights Reserved
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#

#
# WORK IN PROGRESS !
#
# TODO:
# - clean up
# - fix occasional off-by-one error in left-hand side of selection
# - add inversion
# - add deglitching
# - corner case: handle DC traces (sdio-121-A)
#


from Tkinter import *
from tmc.decode import *
from math import ceil, floor
from string import maketrans


channel_y_step = 55
channel_y_offset = 15
channel_y_height = 20


class channel:

    color_normal = "green"
    color_selected = "yellow"
    color_edit = "red"
    edit_border = 2
    decode_color = "white"
    decode_bg_color = "#4040ff"
    decode_font = "-*-helvetica-medium-r-*-*-12-*-*-*-*-*-*-*"
    label_color = "white"
    label_font = "-*-helvetica-medium-r-*-*-10-*-*-*-*-*-*-*"

    def __init__(self, main, number, data, label):
	self.main = main
	self.number = number
	self.tag = "d_%d" % number
	self.zoom_tag = "d_z_%d" % number
	self.d = data
	self.label = label 
	self.edits = []
	self.draw()

    def draw(self):
	if self.main.pos0 > len(self.d):
	    return

	if self.label:
	    self.main.w.create_text(2, self.y(1)-1, anchor = "sw",
	    text = self.label,
	    fill = self.label_color, font = self.label_font,
	    tags = "d_lbl_"+self.tag)
	for pos in self.edits:
	    if pos < self.main.pos0:
		break
	    self.main.w.create_rectangle(
	      self.x(pos)-self.edit_border, self.y(0)+self.edit_border,
	      self.x(pos+1)+self.edit_border, self.y(1)-self.edit_border,
	      fill = self.color_edit, outline = self.color_edit,
	      tags = "d_ed_"+self.tag)

	pos0 = max(0, self.main.pos0)
	last = self.d[pos0]
	line = [ self.x(pos0), self.y(last) ]
	for i in range(pos0, len(self.d)):
	    if self.d[i] != last:
		line.extend((self.x(i), self.y(self.d[i])))
	    x1 = self.x(i+1)
	    if x1 > self.main.xres:
		break
	    if self.d[i] == last and i != pos0:
		line.pop()
		line.pop()
	    line.extend((x1, self.y(self.d[i])))
	    last = self.d[i]
	self.main.w.create_line(
	  fill = self.color_normal, tags = self.tag, *line)

    def draw_zoom(self):
	mag = float(self.main.xres)/self.main.samples
	last = self.d[0]
	line = [ 0, self.zoom_y(last) ]
	for i in range(0, len(self.d)):
	    if self.d[i] != last:
		line.extend((i*mag, self.zoom_y(self.d[i])))
	    elif i:
		line.pop()
		line.pop()
	    line.extend(((i+1)*mag, self.zoom_y(self.d[i])))
	    last = self.d[i]
	self.main.wz.create_line(
	  fill = self.color_normal, tags = self.zoom_tag, *line)

    def redraw_zoom(self):
	self.main.wz.delete(self.zoom_tag)
	self.draw_zoom()
	
    def redraw(self):
	self.main.w.delete(self.tag)
	self.main.w.delete("d_ed_"+self.tag)
	self.main.w.delete("d_lbl_"+self.tag)
	self.draw()

    def x(self, sample):
	return int(round(self.main.x0+(sample-self.main.pos0)*self.main.mag))

    def y(self, value):
	return self.number*channel_y_step+channel_y_offset+channel_y_height- \
	  value*20

    def zoom_y(self, value):
	return self.number*6-value*2+5

    def select(self):
	self.main.w.itemconfig(self.tag, fill = self.color_selected)

    def deselect(self):
	self.main.w.itemconfig(self.tag, fill = self.color_normal)

    def set_decoder(self, decoder, both_edges):
	self.decoder = decoder
	self.both_edges = both_edges

    def begin_decode(self, pos):
	x0 = self.x(pos)
	x1 = x0+self.main.mag
	y = self.y(0)
	self.main.w.create_text(x0+5, y+6,
	  fill = self.decode_color, anchor = "nw", tags = "d_fg_"+self.tag,
	  font = self.decode_font)
	rect = self.main.w.create_rectangle(x0, y+5, x1, y+20,
	  fill = self.decode_bg_color, outline = self.decode_bg_color,
	  tags = "d_bg_"+self.tag)
	self.main.w.tag_raise("d_bg_"+self.tag)
	self.main.w.tag_raise("d_fg_"+self.tag)

    def pre_decode(self, selected, bits):
	global d_table

	if not selected:
	    d_table[self.decoder] = bits

    def decode(self, selected, bits):
	y = self.y(0)
	x0 = self.x(self.main.decode_from)
	x1 = self.x(self.main.cur.pos+1)
	if selected:
	    text = d_counter(bits)
	else:
	    text = self.decoder(bits)
	self.main.w.itemconfig("d_fg_"+self.tag, text = text)
	self.main.w.coords("d_bg_"+self.tag, x0, y+5, x1, y+20)

    def end_decode(self):
	self.main.w.delete("d_fg_"+self.tag)
	self.main.w.delete("d_bg_"+self.tag)

    def edit(self, pos):
	self.d[pos] = 1-self.d[pos]
	if pos in self.edits:
	    self.edits.remove(pos)
	else:
	    self.edits.append(pos)
	self.redraw()


# Base class to have access to the pretty-printing methods.

class measurement_base:

    def pretty_float_common(self, f):
	if f < 0:
	    sign = True
	    f = -f
	else:
	    sign = False

	if f < 1e-6:
	    unit = "n"
	    f *= 1e9
	elif f < 1e-3:
	    unit = "u"
	    f *= 1e6
	elif f < 1:
	    unit = "m"
	    f *= 1e3
	elif f < 1e3:
	    unit = " "
	elif f < 1e6:
	    unit = "k"
	    f *= 1e-3
	elif f < 1e9:
	    unit = "M"
	    f *= 1e-6
	elif f < 1e12:
	    unit = "G"
	    f *= 1e-9
	elif f < 1e15:
	    unit = "T"
	    f *= 1e-12
	else:
	    print f
	    raise hell

	return sign, f, unit

    def pretty_float(self, f):
	sign, f, unit = self.pretty_float_common(f)
	return "%s%7.3f%s" % ((" ", "-")[sign], f, unit)

    #
    # pretty_float_trim is like pretty_float but it removes trailing zeroes,
    # e.g., 5.100 becomes 5.1. While this looks like a nice idea, I found the
    # constantly changing format somewhat irritating, so I don't use it for the
    # continuously updated measurements.
    #
    # It's also inaccurate in the sense that these numbers represent physical
    # units well above quantum level and therefore cannot be exact.
    #

    def pretty_float_trim(self, f):
	sign, f, unit = self.pretty_float_common(f)

	if round(f*1000) == round(f)*1000:
	    dec = 0
	elif round(f*1000) == round(f*10)*100:
	    dec = 1
	elif round(f*1000) == round(f*100)*10:
	    dec = 2
	else:
	    dec = 3

	fract = (0, 1+dec)[dec > 0]

	return "%s%*.*f%*s%s" % (
	  (" ", "-")[sign],
	  fract+3, dec, f,
	  4-fract, "",
	  unit)


class measurement(measurement_base):

    background_color = "black"
    color = "white"
    font = "-*-helvetica-medium-r-*-*-12-*-*-*-*-*-*-*"
    # sigh, where's a mono-spaced font with a sans-serif "1" when you need
    # one ...
    font = "-*-fixed-medium-r-normal-*-14-*-*-*-*-*-iso8859-*"

    def __init__(self, master, t0, dt, prefix = ""):
        self.fn = (
	  ( self.show_time, "s " ),
	  ( self.show_samples, "Sa" ),
	  ( self.show_frequency, "Hz" ))

	self.t0 = t0
	self.dt = dt
	self.prefix = prefix
	self.last_prefix = None
	self.var = StringVar()
	self.button = Button(master, textvariable = self.var, width = 14,
	  relief = FLAT, borderwidth = 0, font = self.font,
          fg = self.color, bg = self.background_color,
	  activeforeground = self.color,
	  activebackground = self.background_color,
	  command = self.button)
	self.button.pack(side = LEFT)
	self.index = 0
	self.last = None
	self.hide()

    def button(self):
	self.index += 1
	self.index %= len(self.fn)
	self.__show(self.last)

    def show_samples(self, samples):
	self.var.set(self.prefix+"%8d Sa" % samples)

    def show_time(self, samples):
	t = samples*self.dt+self.t0
	self.var.set(self.prefix+self.pretty_float(float(t))+"s")

    def show_frequency(self, samples):
	t = samples*self.dt+self.t0
	if abs(t) < 1e-12:
	    self.hide()
	else:
	    self.var.set(self.prefix+self.pretty_float(1.0/t)+"Hz")

    def __show(self, samples):
	if samples is None:
	    self.hide()
	else:
	    self.fn[self.index][0](samples)

    def show(self, samples):
	if samples != self.last or self.prefix != self.last_prefix:
	    self.__show(samples)
	    self.last = samples
	    self.last_prefix = self.prefix

    def hide(self):
	self.var.set(self.prefix+" "*9+self.fn[self.index][1])
	self.last = None

    def remove(self):
	self.button(delete)


class cursor:

    color = "white"

    def __init__(self, main):
	self.main = main
	self.tag = main.w.create_line(0, 0, 0, 0, fill = self.color, width = 1)
	self.zoom_tag = \
	  main.wz.create_line(0, 0, 0, 0, fill = self.color, width = 1)
	self.x = 0
	self.pos = 0

    def set(self):
	x = self.main.ch[0].x(self.pos)
	self.main.w.coords(self.tag, x, 0, x, self.main.yres)
	self.main.w.tag_raise(self.tag)
	x = int(float(self.pos)/self.main.samples*self.main.xres)
	self.main.wz.coords(self.zoom_tag, x, 0, x, self.main.zres)
	self.main.wz.tag_raise(self.zoom_tag)
	self.main.meas_pos.show(self.pos)
	if self.main.user_pos is not None:
	    self.main.meas_user.show(self.pos-self.main.user_pos)

    def x_to_pos(self, x):
	s = self.main.pos0+int(round((x-self.main.x0)/self.main.mag))
	if s < 0:
	    return 0
	if  s >= self.main.samples:
	    return self.main.samples-1
	return s

    def move(self, x):
	s = self.x_to_pos(x)
	self.x = x
	if s != self.pos:
	    self.pos = s
	    self.set()


class main_window:

    background_color = "black"
    selection_color = "#808080"
    data_color = "white"
    zoom_color = "#808080"
    user_marker_color = "red"

    tick_font = "-*-helvetica-medium-r-*-*-10-*-*-*-*-*-*-*",
    tick_color = "#c0c0c0"
    tick_mark_color = "#c07070"

    si = ("p", "n", "u", "m", "", "k", "M", "G", "T")

    def __init__(self, master, d, t0, sample_step, labels):
	channels = len(d)
	self.samples = len(d[0])
	self.t0 = t0
	self.sample_step = sample_step
	self.xres = self.samples
	self.yres = channels*channel_y_step+channel_y_offset-5
	self.zres = channels*6+2
	self.geometry()
	self.w3 = self.control_window(master)
	self.wz = self.zoom_window(master)
	self.w = self.main_window(master)
	self.wt = self.tick_window(master)
	self.wd = self.meas_window(master)
	self.setup_events(master)
	self.decode_from = None
	self.user_pos = None
	self.ch = []
	self.cur = cursor(self)
	for ch in range(0, channels):
	    if labels is None:
		label = None
	    else:
		label = labels[ch]
	    self.ch.append(channel(self, ch, d[ch], label))
	    self.decoder_menu(self.w3, self.ch[-1], ch)
	self.selected = None
	self.update_zoom()

    def geometry(self):
	self.x0 = 0
	self.pos0 = 0
	self.mag = 1.0

    def main_window(self, master):
	w = Canvas(master,
	  width = self.xres,
	  height = self.yres, bg = self.background_color)
	w.pack(expand = 1, fill = "x")
	return w

    def zoom_window(self, master):
	w = Canvas(master,
	  width = self.xres,
	  height = self.zres, bg = self.zoom_color)
	w.pack(expand = 1, fill = "x")
	self.zoom_area = w.create_rectangle(0, 0, 0, 0,
	  fill = self.background_color, width = 0)
	return w

    def meas_window(self, master):
	w = Frame(master, width = self.xres, height = 20,
	  bg = self.background_color)
	w.pack(expand = 1, fill = "x")

	self.meas_start = measurement(w, 0, self.sample_step, "SEL ")
	self.meas_pos = measurement(w, self.t0, self.sample_step, "CUR ")
	self.meas_user = measurement(w, 0, self.sample_step, "USR ")
	self.meas_width = measurement(w, 0, self.sample_step)

	self.meas_status = measurement_base()
	self.status_var = StringVar()
	self.status = Label(w, textvariable = self.status_var,
	  relief = FLAT, borderwidth = 0, font = measurement.font,
          fg = measurement.color, bg = measurement.background_color,
	  activeforeground = measurement.color,
	  activebackground = measurement.background_color)
	self.status.pack(side = RIGHT, fill = "x")
	return w

    def tick_window(self, master):
	w = Canvas(master, width = self.xres, height = 20,
	  bg = self.background_color)
	w.pack(expand = 1, fill = "x")
	return w

    def control_window(self, master):
	w = Frame(master, width = 135, bg = self.background_color)
	w.pack(side = LEFT, fill = BOTH, expand = 1)
	b = Button(master, text = "Quit", relief = FLAT, command = master.quit)
	b.place(x = 3, y = self.yres+16+self.zres+2)
	return w

    def setup_events(self, master):
	self.w.bind("<Motion>", self.move)
	self.w.bind("<Button-1>", self.button)
	self.w.bind("<Button-3>", self.center)
	self.w.bind("<Button-4>", self.zoom_in)
	self.w.bind("<Button-5>", self.zoom_out)
	self.w.bind("<Configure>", self.resize)

	self.wz.bind("<Button-1>", self.zoom_button)
	self.wz.bind("<Button-3>", self.zoom_button)
	self.wz.bind("<Button-4>", self.zoom_in)
	self.wz.bind("<Button-5>", self.zoom_out)

	master.bind("+", self.zoom_in)
	master.bind("=", self.zoom_in)
	master.bind("-", self.zoom_out)
	master.bind(".", self.center)
	master.bind("c", self.center)
	master.bind("e", self.edit)
	master.bind("<space>", self.user_coord)

    def decoder_menu(self, master, ch, n):
	mb = Menubutton(master, direction = RIGHT, text = decoders[0][0],
	  relief = FLAT, width = 16)
	mb.place(x = 3, y = channel_y_step*n+channel_y_offset-3+self.zres+2)
	mb.menu = Menu(mb, tearoff = 0)
	mb["menu"] = mb.menu

	for dec in decoders:

	    # Brilliant hack. Found it described here:
	    # http://infohost.nmt.edu/tcc/help/pubs/tkinter/extra-args.html

	    def handler(ch = ch, decoder = dec[1], both_edges = dec[2],
	      button = mb, label = dec[0]):
		ch.set_decoder(decoder, both_edges)
		button.config(text = label)

	    mb.menu.add_command(label = dec[0], command = handler)
	    ch.set_decoder(decoders[0][1], decoders[0][2])

    def move(self, event):
	self.cur.move(event.x)
	if self.decode_from is not None:
	    self.decode()
	else:
	    self.move_y(event.y)
	self.measure_width()

    def move_y(self, y):
	ny = y/channel_y_step
	if ny >= len(self.ch):
	    return
	if self.selected is not None:
	    if self.selected.number == ny:
		return
	    self.selected.deselect()
	self.selected = self.ch[ny]
	self.selected.select()

    def measure_width(self):
	if self.selected is None:
	    self.meas_width.hide()
	    return
	val = self.selected.d[self.cur.pos]
	pos = self.cur.pos
	i = pos-1
	while i >= 0 and self.selected.d[i] == val:
	    i -= 1
	width = pos-i-1
	i = pos+1
	while i < self.samples and self.selected.d[i] == val:
	    i += 1
	width += i-pos-1+1	# remove last sample, add center
	self.meas_width.prefix = ("L ", "H ")[val]
	self.meas_width.show(width)

    def zoom(self):
	self.pos0 = self.cur.pos-int(self.cur.x/self.mag)
	self.x0 = self.cur.x-int(self.mag*(self.cur.pos-self.pos0-0.5))
	self.redraw()

    def zoom_in(self, event):
	if self.decode_from is not None:
	    return
	if self.mag > 10:
	    return
	self.mag *= 2.0
	self.zoom()

    def zoom_out(self, event):
	if self.decode_from is not None:
	    return
	if self.xres > self.samples*self.mag:
	    if self.reset():
		self.zoom()
	else:
	    self.mag /= 2.0
	    self.zoom()

    def center(self, event):
	if self.decode_from is not None:
	    return
	x = self.cur.x
	self.cur.x = self.xres/2
	self.zoom()
	self.cur.move(x)

    def zoom_button(self, event):
	if self.decode_from is not None:
	    return
	self.cur.pos = int(float(event.x)/self.xres*self.samples+0.49)
	if self.cur.pos >= self.samples:
	    self.cur.pos = self.samples-1
	self.center(event)

    def user_coord(self, event):
	self.w.delete("user")
	if self.user_pos is None:
	    self.user_pos = self.cur.pos
	    self.meas_user.show(0)
	    self.w.create_polygon(0, 0, 0, 0,
	      fill = self.user_marker_color, outline = "", tags = "user")
	    self.move_user()
	else:
	    self.user_pos = None
	    self.meas_user.hide()

    def move_user(self):
	if self.user_pos is not None:
	    x = self.ch[0].x(self.user_pos)
	    self.w.coords("user",
	      x-7, self.yres+1, x, self.yres-7, x+7, self.yres+1)

    def reset(self):
	end = (self.samples-self.pos0)*self.mag
	# Entire waveform is on screen
	if self.pos0 <= 0 and (self.samples-self.pos0)*self.mag <= self.xres:
	    return False

	# No unused space surrounde waveform
	if self.pos0 >= 0 and end >= self.xres:
	    return False

	pos = self.samples-int(self.xres/self.mag)
	if abs(pos-self.pos0) > abs(self.pos0):
	    self.pos0 = 0
	else:
	    self.pos0 = pos
	self.cur.move(self.cur.x)
	return True
	
    def resize(self, event):
	if self.decode_from is not None:
	    self.end_decode()
	self.xres = event.width
	for ch in self.ch:
	    ch.redraw_zoom()
	self.reset()
	self.zoom()

    def redraw(self):
	for ch in self.ch:
	    ch.redraw()
	self.cur.move(self.cur.x)
	self.update_zoom()
	self.move_user()

    def update_zoom(self):
	self.wz.coords(self.zoom_area,
	  int(float(self.pos0)/self.samples*self.xres), 0,
	  int(ceil((self.pos0+self.xres/self.mag)/self.samples*self.xres)),
	  self.zres)

	shown = min(self.samples, self.xres/self.mag)
	if self.pos0 < 0:
	    shown = min(shown, self.xres/self.mag+self.pos0)
	if self.pos0 > 0:
	    shown = min(shown, self.samples-self.pos0)
	self.status_var.set(
	  "%d/%d Sa, %ss/%ss, step %ss, %s%d" % (shown, self.samples,
	  self.meas_status.pretty_float_trim(shown*self.sample_step).
	    translate(maketrans("", ""), " "),
	  self.meas_status.pretty_float_trim(self.samples*self.sample_step).
	    translate(maketrans("", ""), " "),
	  self.meas_status.pretty_float_trim(self.sample_step).
	    translate(maketrans("", ""), " "),
	  ("/", "x")[self.mag > 0.9],
	  (1.0/self.mag, self.mag)[self.mag > 0.9]))

	self.wt.delete("ticks")

	tick_gap = 100
	tick_inc, tick_exp, tick_unit = \
	  self.tick_floor(self.sample_step*tick_gap/self.mag)
	step = tick_inc*10**tick_exp
	t0 = self.t0+self.sample_step*(self.pos0+self.x0/self.mag)
	t1 = t0+self.sample_step*(self.xres-self.x0)/self.mag

	for tick in range(int(t0*10/step+0.5), int(t1*10/step-0.5)+1):
	    t = tick/10.0*step
	    pos = (t-self.t0)/self.sample_step-self.pos0
	    x = self.x0+pos*self.mag
	    self.wt.create_line(x, 0, x, 5+4*((tick % 5) == 0),
	      tags = "ticks", fill = self.tick_mark_color)
	    if (tick % 10) == 0:
		self.wt.create_text(x, 14,
		  text = self.meas_status.pretty_float_trim(tick/10*step).
		    translate(maketrans("", ""), " ")+"s",
#		  text = str(tick/10*tick_inc)+tick_unit+"s",
		  tags = "ticks",
		  font = self.tick_font,
		  fill = self.tick_color)

    def tick_floor(self, step):
	exp = -12
	for unit in self.si:
	    for dec in range(0, 3):
		for i in range(0, 3):
		    if step < (2, 5, 10)[i]*10**(exp+dec):
			return int(round((1, 2, 5)[i]*10**dec)), exp, unit
	    exp += 3
	raise hell

    def button(self, event):
	if self.decode_from is None:
	    self.begin_decode()
	else:
	    self.end_decode()
	
    def begin_decode(self):
	# @@@FIXME: at mag < 1, this we're usually off by a pixel
	x0 = max(1, self.pos0+int((self.cur.x-self.x0)/self.mag))
	x1 = min(self.samples,
	  self.pos0+int((self.cur.x-self.x0+1)/self.mag)+1)
	# print self.cur.pos, x0, x1
	for pos in range(x0, x1):
	    if self.selected.d[pos-1] != self.selected.d[pos]:
		break
	else:
	    return
	self.decode_from = pos
	for ch in self.ch:
	    ch.begin_decode(pos)
	x0 = self.ch[0].x(pos)
	x1 = x0+self.mag
	rect = self.w.create_rectangle(x0, 0, x1, self.yres,
	  fill = self.selection_color, outline = self.selection_color,
	  tags = "selection")
	self.w.tag_lower(rect)
	self.w.tag_raise(self.cur)
	self.decode()

    def decode(self):
	global d_table

	self.meas_start.show(self.cur.pos-self.decode_from)
	if self.cur.pos < self.decode_from:
	    return
	x0 = self.ch[0].x(self.decode_from)
	x1 = self.ch[0].x(self.cur.pos+1)
	self.w.coords("selection", x0, 0, x1, self.yres)

	d_table = {}
	by_ch = {}
	for ch in self.ch:
	    bits = []
	    i = self.decode_from
	    while i <= self.cur.pos:
		if self.selected.d[i-1] != self.selected.d[i] and \
		  (self.selected.d[i] == self.selected.d[self.decode_from] or
		  self.selected.both_edges):
		    if ch.d[i-1] == ch.d[i]:
			bits.append(ch.d[i])
		    else:
			bits.append("X")
		i += 1
	    by_ch[ch] = bits
	    ch.pre_decode(ch is self.selected and not ch.both_edges, bits)

	d_set_table(d_table)

	for ch in self.ch:
	    ch.decode(ch is self.selected and not ch.both_edges, by_ch[ch])

    def end_decode(self):
	for ch in self.ch:
	    ch.end_decode()
	self.w.delete("selection")
	self.decode_from = None
	self.meas_start.hide()

    def edit(self, event):
	if self.selected is None:
	    return
	self.selected.edit(self.cur.pos)
	self.cur.move(self.cur.x)
	

#
# @@@ The interface is a little ugly with waves containing data and labels,
# but dxplore separating the two. Perhaps the future common wrapper will make
# this more palatable.
#

class dxplore:

    def __init__(self, channels, t0, sample_step, title = None, labels = None):
	self.master = Tk()
	if title is not None:
	    self.master.title(title)
	self.master.resizable(1, 0)
	self.main = main_window(self.master, channels, t0, sample_step, labels)
	self.master.bind("q", self.quit)
	mainloop()

    def quit(self, event):
	self.master.quit()