plz_draw

#!/usr/bin/env python
# encoding: utf-8
u"""plz_draw - draws maps of germany displaying postal codes"""

# Created by Maximillian Dornseif on 2010-01-17.
# Copyright (c) 2010, 2013 HUDORA. All rights reserved.

from __future__ import absolute_import

from optparse import OptionParser
import cairo
import re
import sys


debug = False


def colorstring_to_value(colstr):
    if not colstr.startswith('#') or len(colstr) != 4:
        raise RuntimeError("invalid color %r" % colstr)
    r, g, b = [float(int(x, 16)) / 0xf for x in list(colstr[1:])]
    return (r, g, b)


class NiceCtx(cairo.Context):
    defaultBorderColour = (0x7d / 255.0, 0x7d / 255.0, 0x7d / 255.0)
    # we use 1.6 for a very simple "projection"
    projection = 1.6

    def stroke_border(self, border):
        src = self.get_source()
        width = self.get_line_width()
        self.set_source_rgba(*self.defaultBorderColour)
        self.stroke_preserve()
        self.set_source(src)
        self.set_line_width(width - (border * 2))
        self.stroke()

    def init_geoscale(self, minx, xwidth, miny, yheight):
        # wir normalisieren die Karte, dass sie im Koordinatenraum 0.0-1.0 liegt
        # x ist die Ost/West Ausdehnung
        self.minx = minx
        self.miny = miny
        self.xwidth = xwidth
        self.yheight = yheight
        self.geoscalefactor = 1 / max([xwidth, yheight * self.projection])

    def geoscale(self, x, y):
        # wir spiegeln auch direkt in Nord-Süd Richutng
        newx = (x - self.minx) * self.geoscalefactor
        newy = 1 - ((y - self.miny) * self.geoscalefactor) * self.projection
        return newx, newy


class Canvas:
    def __init__(self, width, height, filename):
        self.width, self.height = width, height
        self.needs_save = False
        if filename.endswith('.pdf'):
            self.surface = cairo.PDFSurface(filename, width, height)
        elif filename.endswith('.svg'):
            # the generated SVG seems broken
            self.surface = cairo.SVGSurface(filename, width, height)
        elif filename.endswith('.ps'):
            self.surface = cairo.PSSurface(filename, width, height)
        elif filename.endswith('.eps'):
            self.surface = cairo.PSSurface(filename, width, height)
            self.surface.set_eps(True)
        else:
            self.surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)
            self.needs_save = filename

        self.background(1, 1, 1)

    def ctx(self):
        ctx = NiceCtx(self.surface)
        ctx.set_line_cap(cairo.LINE_CAP_ROUND)
        ctx.set_line_join(cairo.LINE_JOIN_ROUND)
        ctx.set_antialias(cairo.ANTIALIAS_DEFAULT)
        # Basic scaling so that the context is axexctly 1 unit in the smaller dimension
        # and a little bit more (depending on the aspect ratio) in the other
        #self.ctxscale = min([self.width, self.height])
        # hack to better fit germany - remove the next line if you get stupid results
        self.ctxscale = max([self.width, self.height])
        ctx.scale(self.ctxscale, self.ctxscale)
        self.ctxscale = 1 / float(self.ctxscale)
        return ctx

    def background(self, r, g, b):
        c = self.ctx()
        c.set_source_rgb(r, g, b)
        c.rectangle(0, 0, self.width, self.height)
        c.fill()
        c.stroke()

    def save(self):
        # image background
        ctx = self.ctx()
        ctx.fill()
        # ctx.show_page()
        if self.needs_save:
            self.surface.write_to_png(self.needs_save)
        else:
            self.surface.finish()


def find_dimensions(geoitems, deutschgrenzen):
    # find dimenions of data set
    x = []
    y = []
    #for plz, (lon, lat, name) in geoitems:
    #    x.append(lon)
    #    y.append(lat)
    for track in deutschgrenzen:
        for lon, lat in track:
            x.append(lon)
            y.append(lat)
    return min(x), max(x) - min(x), min(y), max(y) - min(y)


def set_clipping(ctx, deutschgrenzen, borderskip):
    # use borders as clipping region
    for track in deutschgrenzen:
        ctx.move_to(*ctx.geoscale(*track[0]))
        for lon, lat in track[1::borderskip]:
            ctx.line_to(*ctx.geoscale(lon, lat))
        ctx.close_path()
    ctx.clip()


def draw_frontier(ctx, deutschgrenzen, borderskip):
    # draw borders
    for track in deutschgrenzen:
        ctx.move_to(*ctx.geoscale(*track[0]))
        for lon, lat in track[1::borderskip]:  # Borders are too detailed, only use 20% of data
            ctx.line_to(*ctx.geoscale(lon, lat))
        ctx.close_path()
    ctx.stroke()


def should_draw(plz, limit):
    """Determines if data for this PLZ should be drawn (see -l paraleter)"""
    if not limit:
        return True
    else:
        for prefix in limit:
            if plz.startswith(prefix):
                return True
        return False


def get_centercolor(plz, options):
    if plz in options.centercolors:
        return options.centercolors[plz]
    if plz[:4] in options.centercolors:
        return options.centercolors[plz[:4]]
    if plz[:3] in options.centercolors:
        return options.centercolors[plz[:3]]
    if plz[:2] in options.centercolors:
        return options.centercolors[plz[:2]]
    if plz[:1] in options.centercolors:
        return options.centercolors[plz[:1]]
    return None


def draw_centers(ctx, geoitems, options):
    """draw centers of PLZ areas"""

    ctx.set_source_rgb(*colorstring_to_value(options.centerdefaultcolor))
    for plz, (lon, lat, name) in sorted(geoitems):
        if should_draw(plz, options.limit):
            centercolor = get_centercolor(plz, options)
            if centercolor:
                ctx.stroke()
                ctx.set_source_rgb(*centercolor)
                ctx.move_to(*ctx.geoscale(lon, lat))
                ctx.close_path()
                ctx.stroke()
                ctx.set_source_rgb(*colorstring_to_value(options.centerdefaultcolor))
            else:
                ctx.move_to(*ctx.geoscale(lon, lat))
                ctx.close_path()
    ctx.stroke()


class mySite(object):
    def __init__(self, lon, lat, plz):
        self.x = lon
        self.y = lat
        self.plz = plz

    def __repr__(self):
        return "<Site %s %s %s>" % (self.plz, self.lat, self.lon)


def get_areacolor(plz, options):
    if plz in options.areacolors:
        return options.areacolors[plz]
    if plz[:4] in options.areacolors:
        return options.areacolors[plz[:4]]
    if plz[:3] in options.areacolors:
        return options.areacolors[plz[:3]]
    if plz[:2] in options.areacolors:
        return options.areacolors[plz[:2]]
    if plz[:1] in options.areacolors:
        return options.areacolors[plz[:1]]
    return None


def draw_voronoi(ctx, geoitems, options):
    # calculate voronoi diagram for plz areas
    import pygeodb.voronoi
    pts = []
    for plz, (lon, lat, name) in geoitems:
        if should_draw(plz, options.limit):
            pts.append(mySite(lon, lat, plz))
    points, lines, edges = pygeodb.voronoi.computeVoronoiDiagram(pts)

    plzlines = {}
    for (l, p1, p2) in edges:
        # l ist der Index der Original Daten in lines[]
        # p1 und p2 sind der Index der Endpunkte in points[p2]
        x1 = y1 = x2 = y2 = None  # Endpunkte der zu zeichnenen Linie
        a, b, c, plz1, plz2 = lines[l]
        if p1 > -1:
            x1, y1 = points[p1]
        else:
            # pointing to infinity
            x2, y2 = points[p2]
            x1 = x2 - 0.25
            y1 = -1 * ((a * x1 - c) / (b + 0.0000001))
        if p2 > -1:
            x2, y2 = points[p2]
        else:
            # pointing to infinity
            x2 = x1 + 0.25
            y2 = -1 * ((a * x2 - c) / (b + 0.0000001))
        if x1 and y1 and x2 and y2:
            if p1 > -1 and p2 > -1:
                # save for polygon use if not infinity
                plzlines.setdefault(plz1, []).append(((x1, y1), (x2, y2)))
                plzlines.setdefault(plz2, []).append(((x1, y1), (x2, y2)))

            if options.borders:
                # actual drawing
                ctx.move_to(*ctx.geoscale(x1, y1))
                ctx.line_to(*ctx.geoscale(x2, y2))
                ctx.stroke()

    if not options.areas and not options.areacolors:
        return  # we are finished

    # calculate per PLZ polygons
    for plz, lines in plzlines.items():
        polygon = list(lines.pop())
        lenpolygon = len(polygon)
        while lines:
            # clockwise
            needle = polygon[-1]
            for a, b in lines:
                if a == needle:
                    polygon.append(b)
                    del(lines[lines.index((a, b))])
                if b == needle:
                    polygon.append(a)
                    del(lines[lines.index((a, b))])
            # anticlockwise
            needle = polygon[0]
            for a, b in lines:
                if a == needle:
                    polygon.insert(0, b)
                    del(lines[lines.index((a, b))])
                if b == needle:
                    polygon.insert(0, a)
                    del(lines[lines.index((a, b))])

            if lenpolygon == len(polygon):
                # no progress
                break
            lenpolygon = len(polygon)

        # draw the beast
        # TODO: default
        ctx.set_source_rgba(0.95, 0.95, 0.95)
        if len(polygon) > 2:
            areacolor = get_areacolor(plz, options)
            if areacolor:
                ctx.set_source_rgb(*areacolor)
                ctx.move_to(*ctx.geoscale(*polygon[0]))
                for point in polygon[1:]:
                    ctx.line_to(*ctx.geoscale(*point))
                ctx.close_path()
                ctx.fill()
                ctx.set_source_rgba(0.5, 0.5, 0.5)
            elif options.areas:
                ctx.move_to(*ctx.geoscale(*polygon[0]))
                for point in polygon[1:]:
                    ctx.line_to(*ctx.geoscale(*point))
                ctx.close_path()
                ctx.fill()


def draw_citymarker(ctx, name, lon, lat, ctxscale):
    """Draw a Black-White-Red circle and the city name"""
    # draw text
    ctx.set_source_rgb(0, 0, 0)
    xbearing, ybearing, width, height, xadvance, yadvance = ctx.text_extents(name)
    x, y = ctx.geoscale(lon, lat)
    ctx.move_to(x - width / 2, y + (ybearing * 0.8))
    ctx.show_text(name)
    # draw dot
    ctx.set_source_rgb(0.0, 0.0, 0.0)
    ctx.set_line_width(6 * ctxscale)
    ctx.move_to(*ctx.geoscale(lon, lat))
    ctx.close_path()
    ctx.stroke()
    ctx.set_source_rgb(1.0, 1.0, 1.0)
    ctx.set_line_width(5 * ctxscale)
    ctx.move_to(*ctx.geoscale(lon, lat))
    ctx.close_path()
    ctx.stroke()
    ctx.set_source_rgb(1.0, 0.0, 0.0)
    ctx.set_line_width(4 * ctxscale)
    ctx.move_to(*ctx.geoscale(lon, lat))
    ctx.close_path()
    ctx.stroke()


def parse_commandline():
    """Parse the commandline and return information."""

    parser = OptionParser()
    parser.description = __doc__

    parser.set_usage('usage: %prog [options] outputfile. Try %prog --help for details.')
    parser.add_option('--width', type='int', default=768,
                      help='Width of generated Image [%default]')
    parser.add_option('--height', type='int', default=1024,
                      help='Heigth of generated Image [%default]')
    parser.add_option('-l', '--limit', action='append',
                      help='Limit PLZ matiching to this prefifx, can be given more than once.')
    parser.add_option('--nofrontier', action='store_true',
                      help="Don't draw german Frontier")
    parser.add_option('-b', '--borders', action='store_true',
                      help='Draw german Frontier')
    parser.add_option('-m', '--mark', action='append',
                      help='Draw City names of this PLZ.')
    parser.add_option('-c', '--center', type='int', default=50,
                      help='Draw centers of postcode areas that specific size, 0 to disable [%default]')
    parser.add_option('--centerdefaultcolor', type='string', default='#000',
                      help='Draw centers in this color [%default]')
    parser.add_option('--cencol', action='append',
                      help='set color for one or many PLZ centers, e.g. --cencol=424:#f00')
    parser.add_option('-a', '--areas', action='store_true',
                      help='Fill postal areas')
    parser.add_option('--acol', action='append',
                      help='set color for one or many PLZ areas, e.g. --acol=423:#f00')
    parser.add_option('-r', '--read', type='string',
                      help="read list of PLZ to highlight from file")
    parser.add_option('--digits', type='int', default=5,
                      help='Significant digits per PLZ [%default]')
    parser.add_option('--noclip', action='store_true',
                      help="Don't clip around the Border")
    parser.add_option('-d', '--debug', action='store_true', dest='debug',
                      help='Enables debugging mode')

    options, args = parser.parse_args()
    if len(args) != 1:  # we expect no non option arguments
        parser.error('incorrect number of arguments')
    return options, args


def main(options, args):
    """This implements the actual program functionality"""
    global debug, values, stepsize

    import pygeodb
    import pygeodb.borderdata
    geoitems = pygeodb.geodata['DE'].items()

    debug = options.debug
    options.centercolors = {}
    options.areacolors = {}

    if options.cencol:
        # area center (city) colors
        for value in options.cencol:
            prefix, color = value.split(':')
            options.centercolors[prefix] = colorstring_to_value(color)
    if options.acol:
        # area colors
        for value in options.acol:
            prefix, color = value.split(':')
            options.areacolors[prefix] = colorstring_to_value(color)

    values = {}
    maxvalue = 0
    if options.read:
        fd = open(options.read, 'r')
        for line in (x.strip() for x in fd):
            plz = re.sub('[^0-9]+', '', line)
            if plz.startswith('#'):
                continue
            values[plz[:options.digits]] = values.get(plz[:options.digits], 0) + 1
            maxvalue = max([maxvalue, values[plz[:options.digits]]])
        stepsize = 1.0 / maxvalue
        for plzprefix, value in values.items():
            # options.centercolors[plzprefix] = (1-(value*stepsize*0.8), 1-(value*stepsize*0.8), 1)
            options.areacolors[plzprefix] = (1 - (value * stepsize * 0.8), 1 - (value * stepsize * 0.8), 1)

    c = Canvas(options.width, options.height, args[0])
    ctx = c.ctx()

    ctx.init_geoscale(*find_dimensions(geoitems, pygeodb.borderdata.deutschgrenzen))
    ctx.set_line_width(0.1 * c.ctxscale)
    borderskip = 10  # Borders are too detailed, only use 10% of data
    if not options.noclip:
        set_clipping(ctx, pygeodb.borderdata.deutschgrenzen, borderskip)
    if options.borders or options.areas or options.acol:
        draw_voronoi(ctx, geoitems, options)

    ctx.set_line_width(0.5 * c.ctxscale)
    ctx.set_source_rgba(0, 0, 0, 1)
    if not options.nofrontier:
        draw_frontier(ctx, pygeodb.borderdata.deutschgrenzen, borderskip)
    if options.center:
        ctx.set_line_width((options.center / 100.0) * c.ctxscale)
        draw_centers(ctx, geoitems, options)

    if options.mark:
        ctx.reset_clip()
        ctx.set_font_size(0.025)
        for markname in options.mark:
            # find all locations with that city name
            markname = markname.decode('utf-8')
            locationsx = []
            locationsy = []
            for plz, (lon, lat, name) in geoitems:
                if name == markname:
                    locationsx.append(lon)
                    locationsy.append(lat)
            if not locationsx:
                sys.stderr.write("%r not found\n" % markname)
            else:
                # average
                draw_citymarker(ctx, markname, sum(locationsx) / len(locationsx), sum(locationsy) / len(locationsy), c.ctxscale)
    c.save()


if __name__ == '__main__':
    main(*parse_commandline())