treedraw.py

##    A simple program that showcases tree drawing by fractals
##    Copyright (C) 2016  Mario Abarca
##
##    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 3 of the
##    License, or (at your option) any later version.
##
##    This program is distributed in the hope that it will be useful,
##    but WITHOUT ANY WARRANTY; without even the implied warranty of
##    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
##    General Public License for more details.
##
##    You should have received a copy of the GNU General Public License
##    along with this program.  If not, see
##    <http://www.gnu.org/licenses/>.

import random
import tkinter as tk
from collections import namedtuple
from math import sin, cos, pi, sqrt, atan


## A Point in the plane is given by its X and Y coordinates.
class Point2D(namedtuple('Point2D', ['x', 'y'])):
    ## We may add two points by adding its coordinate values component-wise.
    def __add__(self, other: 'Point2D'):
        return Point2D(self.x + other.x, self.y + other.y)

    ## Same for substraction.
    def __sub__(self, other: 'Point2D'):
        return Point2D(self.x - other.x, self.y - other.y)

    ## We may also scale a point by multiplying its X and Y values by
    ## some number.
    def __mul__(self, scale: float):
        assert isinstance(scale, float)
        return Point2D(scale*self.x, scale*self.y)

    __rmul__ = __mul__


    ## Points can also be represented by an angle and a distance from
    ## the origin. This is the so-called polar-form.
    @staticmethod
    def from_polar(radius, argument):
        return radius*Point2D(cos(argument), sin(argument))

    def to_polar(self):
        return sqrt(self.x**2 + self.y**2), atan(self.x/self.y)


## A color in a computer screen can be represented by giving the
## intensities of the red, green and blue LED's
class Color(namedtuple('Color', ['red', 'green', 'blue'])):
    ## We may want to mix two colors in different proportions. For
    ## example, to mix Yellow and Green we need to specify a ratio (a
    ## number between 0.0 and 1.0) which tell us how much closer to
    ## Green we want the color to be; i.e. a ratio=0.0 is all Yellow, a
    ## ratio=1.0 is all Green and ratio=0.5 is half yellow, half green.
    def mix_with(self, other, ratio):
        def convex_combination(x, y, r):
            return x * (1 - r) + y * r

        assert 0.0 <= ratio <= 1.0
        red = convex_combination(self.red, other.red, ratio)
        green = convex_combination(self.green, other.green, ratio)
        blue = convex_combination(self.blue, other.blue, ratio)
        return Color(red, green, blue)

    ## Computer colors are usually represented in hexadecimal.
    def hex(self):
        red_byte = '{0:0>2x}'.format(round(self.red * 255))
        green_byte = '{0:0>2x}'.format(round(self.green * 255))
        blue_byte = '{0:0>2x}'.format(round(self.blue * 255))
        return ''.join(['#', red_byte, green_byte, blue_byte])


## These are the default colors for leaves and trunks respectively:
FOREST_GREEN = Color(red=0.133, green=0.545, blue=0.133)
SADDLE_BROWN = Color(red=0.545, green=0.271, blue=0.075)

## A branch is represented by 4 values:
##
## + The height respective to the length of the trunk. 0.0 means that
## the branch begin at the base and 1.0 means that it begin at the top
##
## + The angle respective to the angle of the trunk (initialy, the trunk
## is drawn at 90° or pi/2 radians).
##
## + The length of the branch respective to the length of the trunk;
## i.e. a branch with length 1.0 will have the same length as the trunk
## itself.
##
## + The width of the branch respective to the width of the trunk.
Branch = namedtuple('Branch', ['height', 'angle', 'length', 'width'])


## A fractal tree is basically a collection of branches (as defined
## above) together with some drawing information, such as:
##
## + The depth of the tree, which is the number of stright lines you
## need to connect the root with any leave (a perfect fractal would have
## this number equal to infinity, but that is not the case in real
## life).
##
## + Trunk and leaf colors; all the intermediate strokes will be drawn
## with a collor in between those two.
##
## + Trunk length and width.
class FractalTree:
    def __init__(self, **kwargs):
        self.depth = kwargs.get('depth', 8)
        self.branches = []
        for branch in kwargs.get('branches', []):
            self.add_branch(branch)
        self.trunk_width = kwargs.get('trunk_width', 16)
        self.trunk_length = kwargs.get('trunk_length', 128)
        self.trunk_color = kwargs.get('trunk_color', SADDLE_BROWN)
        self.leaf_color = kwargs.get('leaf_color', FOREST_GREEN)

    def add_branch(self, branch: Branch):
        assert isinstance(branch, Branch)
        self.branches.append(branch)

    ## Here is the drawing method of a tree. We just need to know on
    ## which canvas will be drawed, and at which coordinates.
    def draw(self, canvas: tk.Canvas, x: float, y: float):
        ## This method is recursive: it draws an stroke at a given
        ## position, angle, width, and length, and then it recursively
        ## draws the strokes of its branches in an smaller scale.
        def stroke(x1, y1, angle, width, length, level):
            if level > self.depth:      # Did we reach a leaf?
                return                  # stop growing this branch
            ## Compute the ending point of the current stroke:
            x2, y2 = x1 + length*cos(angle), y1 + length*sin(angle)
            ## Compute the color of the current stroke:
            ratio = level/self.depth
            color = (self.trunk_color.mix_with(self.leaf_color, ratio)).hex()
            ## Draw the stroke:
            canvas.create_line(x1, y1, x2, y2,
                               fill=color, width=width, capstyle=tk.ROUND)
            p = Point2D(x1, y1)
            q = Point2D(x2, y2)
            for branch in self.branches:    # For each branch:
                ## ...compute its position, angle, width, and length
                ## respective to the current one
                new_x, new_y = p + (q - p)*branch.height
                new_angle = angle + branch.angle
                new_width = width * branch.width
                new_length = length * branch.length
                new_level = level + 1
                ## Recursively draw each branch:
                stroke(new_x, new_y, new_angle,
                       new_width, new_length, new_level)
        ## This is the initial call: draw a stroke with the length and
        ## width of the trunk at 90° at the given x, y coordinates.
        stroke(x, y, -pi/2, self.trunk_width, self.trunk_length, 0)


## The Fractal Tree Application just creates a window with a single
## canvas wich reacts to the user input. Currently, a left-click will
## draw a tree at the given position and a right-click will erase the
## canvas.
class FractalTreeApp(tk.Tk):
    def __init__(self):
        super().__init__()
        self.rowconfigure(0, weight=1)
        self.columnconfigure(0, weight=1)
        self.canvas = tk.Canvas(master=self)
        self.canvas.configure(background='white')
        self.canvas.grid(sticky=tk.NSEW)
        self.canvas.bind('<Button-1>', self.on_left_click)
        self.canvas.bind('<Button-3>', self.on_right_click)

    def on_left_click(self, event):
        x, y = event.x, event.y
        # Create a random tree
        tree = FractalTree()
        h = 0.1*y
        tree.trunk_length = 4*h
        tree.trunk_width = 0.75*h
        branching_factor = random.randint(2, 4)
        for k in range(branching_factor):
            height = random.uniform(0.33, 1.0)
            length = random.uniform(0.5, 0.75)
            angle = random.uniform(-pi/3, pi/3)
            width = length*0.75
            tree.add_branch(Branch(height, angle, length, width))
        tree.draw(self.canvas, x, y)

    def on_right_click(self, event):
        self.canvas.delete('all')

if __name__ == '__main__':
    print('Click anywhere on the new window to plant a seed...')
    FractalTreeApp().mainloop()