pyramid.py

import svgwrite
import math
import json

class Pyramid:
    def __init__(self, data, filename, height=500):
        self.data = data
        self.filename = filename
        self.height = height
        self.width = 2 * height
        self.drawing = svgwrite.Drawing(filename, height=self.height, width=self.width)

    def draw_polygon(self, points, color="#00ffff"):
        shape = svgwrite.shapes.Polygon(
                points,
                stroke=svgwrite.rgb(0, 0, 0, '%'),
                stroke_width=1,
                stroke_opacity=100,
                fill=color,
                fill_opacity=100)

        self.drawing.add(shape)

    def add_text(self, txt, coord, fontsize=30):
        g = self.drawing.g(
                style="font-size:%s;font-family:DejaVu Sans;text-align:center" % fontsize
                )

        g.add(self.drawing.text(txt, insert=coord))
        self.drawing.add(g)

    def draw_pyramid(self, colour="#ff0000"):
        midpoint = self.width / 2
        points=[
                (midpoint, 0),
                (self.width, self.height),
                (0, self.height),
                ]
        shape = svgwrite.shapes.Polygon(
                points,
                stroke=svgwrite.rgb(0, 0, 0, '%'),
                stroke_width=1,
                stroke_opacity=100,
                fill=colour,
                fill_opacity=100)
        self.drawing.add(shape)


    def volume_at(self, base, height):
        squareside = (base - height)
        if (squareside > 0):
            rectangle = (squareside * 2) * height
        else:
            rectangle = 0

        remainder = height * height
        # print("(total: %s) - volume at %s, %s = %s" % ( (HEIGHT*WIDTH) / 2, base, height, rectangle + remainder))
        return rectangle + remainder


    def calculate_heights(self, arr):
        volume = (self.height * self.width) / 2
        total = sum(arr)
        unit = volume / total
        volumes = [(unit * x) for x in arr]
        print("volume %s, total %s, unit %s, volumes : %s" % (volume, total, unit, volumes))

        heights = []
        base = self.height
        height = 1

        for v in volumes:
            running = True
            while running:
                if (self.volume_at(base, height) >= v):
                    heights.append([base, height])
                    base -= height
                    height = 1
                    running = False
                elif (height > self.height):
                    heights.append([base, height])
                    base -= height
                    height = 1
                    running = False
                else:
                    height +=1

        return heights

    def points_for(self, coor):
        base = coor[0]
        height = coor[1]
        midpoint = self.width/2
        bottom = base
        top = (base - height)

        return [
                (midpoint, base),
                ((midpoint - bottom), base),
                ((midpoint - top), (base - height)),
                ((midpoint + top), (base - height)),
                ((midpoint + bottom), base)
                ];

    def add_segments(self):
        coords = self.calculate_heights([item["number"] for item in data["items"]])
        for idx, coord in enumerate(coords):
            item = self.data["items"][idx]
            self.draw_polygon(
                    self.points_for(coord),
                    color=item["color"],
                    )
            label = item["label"]
            labelshift = (len(label) / 2) * 15
            self.add_text(
                    txt=item["label"],
                    coord=(
                        (self.width / 2) - labelshift,
                        coord[0] - (coord[1]/2) + 12)
                    )

            num = "%s" % item["number"]
            numshift = (len(num) / 2) * 15
            self.add_text(
                    txt=num,
                    coord=(
                        (self.width / 2) - numshift,
                        coord[0] - (coord[1]/2) + 36),
                    fontsize=16
                    )

    def save(self):
        self.drawing.save()

with open('sample.json', 'r') as json_file:
    data = json.load(json_file)
    pyramid = Pyramid(data, 'sample.svg', 600)
    pyramid.add_segments()
    pyramid.save()