#!/usr/bin/python
import pygtk; pygtk.require('2.0')
import gtk
import gobject
import cairo
import colorsys
import random
import math

class Trace(object):
    def __init__(self, data=[], xdiv=.001, ydiv=1, yoff=0, color=(0,1,1)):
        self.data = data
        self.xdiv = xdiv
        self.ydiv = ydiv
        self.yoff = yoff
        self.color = color
        self.cache = None
        self.selected = False

    def update_cache(self, width):
        if self.cache and width == self.cache_width:
            return self.cache
        self.cache_width = width
        samples_per_pixel = 10. / self.xdiv / width

        if samples_per_pixel > 1:
            self.sparse = False
            self.update_cache_dense(samples_per_pixel)
        else:
            self.sparse = True
            self.cache = [(i/samples_per_pixel, v)
                            for i, v in enumerate(self.data)]

    def update_cache_dense(self, samples_per_pixel):
        cache = []

        def stats(s):
            s.sort()
            return s[len(s)/2], s[0], s[-1] 

        def _putcache(i, p):
            while len(self.cache) <= i:
                self.cache.append([])
            self.cache[i].append(p)

        for i, v in enumerate(self.data):
            ci = int(i/samples_per_pixel)
            while len(cache) <= ci:
                cache.append([])
            cache[ci].append(v)

        self.cache = [stats(s) for s in cache]

    def kill_cache(self):
        self.cache = None

    def _cacheprop(attr, doc=None):
        _attr = "_" + attr
        def _get(self):
            return getattr(self, _attr)

        def _set(self, value):
            self.kill_cache()
            setattr(self, _attr, value)

        def _del(self, value):
            self.kill_cache()
            return delattr(self, _attr)
        return property(_get, _set, doc=doc)

    data = _cacheprop('data')
    xdiv = _cacheprop('xdiv')
    dt = _cacheprop('dt')

    def draw(self, canvas, xo, width, height):
        self.update_cache(width)

        canvas.save()

        scale = height / self.ydiv / 10
        yoff = -self.yoff
        def Y(y): return height-(y+yoff)*scale
        canvas.set_line_width(1)
        try:
            if self.sparse: return self.draw_sparse(canvas, xo, width, Y)
            else: return self.draw_dense(canvas, xo, width, Y)
        finally:
            canvas.restore()
       

    def draw_dense(self, canvas, xo, width, Y):
        cache = self.cache
        if xo > 0: r = range(width)
        else: r = range(-xo, width)
#        for row in r:
#            crow = row + xo
#            if crow >= len(cache): break
#            data = cache[crow]
#            canvas.rectangle(row, data[1], row+1, data[2])
#        self.set_color(canvas, .25)
#        canvas.fill()

        self.set_color(canvas, .80)
        for row in r:
            crow = row + xo
            if crow >= len(cache): break
            data = cache[crow]
            canvas.move_to(row, Y(data[1]))
            canvas.line_to(row, Y(data[2]))
        canvas.stroke()

        self.set_color(canvas, .80)
        for row in r:
            crow = row + xo
            if crow >= len(cache): break
            data = cache[crow]
            if row == r[0]:
                canvas.move_to(row, Y(data[0]))
            else:
                canvas.line_to(row, Y(data[0]))
        canvas.stroke()

    def draw_sparse(self, canvas, xo, width, Y):
        cache = self.cache
        first = True

        self.set_color(canvas, .80)
        for x, y in cache:
            if x < xo: continue
            if x > xo+width: break

            if first:
                first = False
                canvas.move_to(x, Y(y))
            else:
                canvas.line_to(x, Y(y))
        canvas.stroke()

    def set_color(self, canvas, alpha):
        r, g, b = self.color
        canvas.set_source_rgba(r*alpha, g*alpha, b*alpha, alpha)


def draw_reticle(canvas, width, height):
    d = .01 * min(width, height)
    canvas.set_line_width(2)
    for row in range(11):
        for col in range(11):
            x = width * col / 10.
            y = height * row / 10.
            if row == 5 or col == 5:
                canvas.set_source_rgba(.8, .8, .8)
            else:
                canvas.set_source_rgba(.3, .3, .3)
            canvas.move_to(x-d, y)
            canvas.line_to(x+d, y)
            canvas.move_to(x, y-d)
            canvas.line_to(x, y+d)
            canvas.stroke()

    canvas.set_source_rgba(.1, .1, .1)
    d = d / 3.
    if height > 500:
        for col in range(11):
            for row in range(101):
                x = width * col / 10.
                y = height * row / 100.
                canvas.move_to(x-d, y)
                canvas.line_to(x+d, y)

    canvas.set_line_width(1)
    if width > 500:
        for col in range(101):
            for row in range(11):
                x = width * col / 100.
                y = height * row / 10.
                canvas.move_to(x, y-d)
                canvas.line_to(x, y+d)
    canvas.stroke()

def draw_traces_to_surface(surface, traces):
    canvas = cairo.Context(surface)
    draw_traces_to_canvas(canvas, traces, surface.get_width(), surface.get_height())

def draw_traces_to_canvas(canvas, traces, width, height):
    canvas.set_source_rgb(0,0,0)
    canvas.set_operator(cairo.OPERATOR_SOURCE)
    canvas.paint()

    canvas.set_operator(cairo.OPERATOR_ADD)
    draw_reticle(canvas, width, height)
    for t in traces:
        t.draw(canvas, 0, width, height)

def draw(traces, target):
    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, 1024, 1024)
    draw_traces_to_surface(surface, traces)
    surface.write_to_png(target)

def f1(t):
    return math.sin(t)

def f2(t):
    return random.gauss(0, .25)

def f3(t):
    return f1(t) + f1(3*t) + f1(4*t)

def f4(t):
    return f2(t) + f3(t)

DT_DEF = .002
def make_trace(fn, n=10000, dt=DT_DEF):
    return [fn(dt*i) for i in range(n)]


d = make_trace(f4)
e = make_trace(lambda t: f3(2*t))
t = Trace(d)
u = Trace(e, color = (1,0,1))

t.yoff = -5
u.yoff = -3

traces = [t,u]
for n, xdiv in enumerate((.02, .005, .001)):
    t.xdiv = xdiv
    u.xdiv = xdiv
    draw(traces, "trace%d.png" % n)

t.xdiv = .01
u.xdiv = .01

# Create a GTK+ widget on which we will draw using Cairo
class Screen(gtk.DrawingArea):

    # Draw in response to an expose-event
    __gsignals__ = { "expose-event": "override" }

    # Handle the expose-event by drawing
    def do_expose_event(self, event):

        # Create the cairo context
        cr = self.window.cairo_create()

        # Restrict Cairo to the exposed area; avoid extra work
        cr.rectangle(event.area.x, event.area.y,
                event.area.width, event.area.height)
        cr.clip()

        self.draw(cr, *self.window.get_size())

    def draw(self, cr, width, height):
        draw_traces_to_canvas(cr, traces, width, height)

# GTK mumbo-jumbo to show the widget in a window and quit when it's closed
def run(Widget):
    window = gtk.Window()
    window.connect("delete-event", gtk.main_quit)
    widget = Widget()
    widget.show()
    window.add(widget)
    window.present()
    gtk.main()

run(Screen)