import pygame
import time
import sys
import random

WIDTH, HEIGHT = 900, 600
WHITE = (255, 255, 255)
BLUE = (0,100,200)
PINK = (255, 0, 120)

#### SETTINGS ###

LINE_WIDTH = 1
TREE_ROOT = [WIDTH/2, HEIGHT]
ROOT_LEN = 150
TREE_LEN = 10
ANGLE = 30
SHORTEN = 0.75

#### SETTINGS ###

pygame.init()
canvas = pygame.display.set_mode((WIDTH,HEIGHT))
background  = pygame.Surface((WIDTH,HEIGHT))
font = pygame.font.SysFont("Calibri", 30)
pygame.display.set_caption('Pretty tree')
clock = pygame.time.Clock()

def quitGame():
    pygame.quit()
    sys.exit(0)

class Branch():
    def __init__(self, start, end):
        self.start = start
        self.end = end
        self.shown = False

    def update(self):
        pass

    def show(self):
        pygame.draw.line(canvas, PINK, self.start, self.end, LINE_WIDTH)

    def newBranch(self, angle, offset):
        new = self.end - self.start
        new = new.rotate(angle)
        newEnd = new * SHORTEN + offset
        return (newEnd)

    def wind(self):
        windX = random.randint(-1,1)
        self.end.x += windX
        windY = random.randint(-1,1)
        self.end.y += windY

start = pygame.math.Vector2(WIDTH/2, HEIGHT)
end = pygame.math.Vector2(WIDTH/2, (HEIGHT) - ROOT_LEN)
tree = []
tree.append(Branch(start,end))

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            quitGame()
        if event.type == pygame.MOUSEBUTTONUP:
            for i in range(len(tree)):
                offset = tree[i].end

                if tree[i].shown != True:
                    angle = random.randrange(10, 60, 10)
                    tree.append(Branch(tree[i].end, tree[i].newBranch(angle, offset) ))

                    angle = random.randrange(-60, -10, 10)
                    tree.append(Branch(tree[i].end, tree[i].newBranch(angle, offset) ))

                tree[i].shown = True

    for i in range(len(tree)):
        tree[i].show()
        # tree[i].wind()

    pygame.display.flip()
    clock.tick(60)