from pgl import GWindow, GOval, GRect
import random

GW_WIDTH = 800
GW_HEIGHT = 800

def random_color():
    color = "#"
    for _ in range(6):
        color += random.choice("0123456789ABCDEF")
    return color

class Firework:
    """ Creates a new firework with initial flight and then 
    explosion. 
    """
    def __init__(self, size):
        self.obj = GOval(GW_WIDTH/2, GW_HEIGHT, size, size)
        self.obj.set_filled(True)
        self.obj.set_color("white")
        self.speed = 5
        self.heading = random.randint(60,120)
        self.fuse = random.randint(50,100)
        self.maxsize = random.randint(60,100)
        self.color = random_color()
        self.mode = 0

    def get_object(self):
        """ Returns the firework graphical object. """
        return self.obj

    def should_terminate(self):
        """ Checks if the firework should be removed. """
        return self.mode > 1

    def move(self):
        """ Moves the firework in its initial flight. """
        self.obj.move_polar(self.speed, self.heading)
        self.fuse -= 1
        if self.fuse < 0:
            self.mode += 1
            self.obj.set_color(self.color)

    def explode(self):
        """ Grows the firework explosion upon detonation. """
        R = 2
        x = self.obj.get_x()
        y = self.obj.get_y()
        S = self.obj.get_width()
        self.obj.set_bounds(x-R/2, y-R/2, S+R, S+R)
        if self.obj.get_width() >= self.maxsize:
            self.mode += 1

    def update(self):
        """ Controls what the firework should be doing during 
        each stage. 
        """
        if self.mode == 0:
            self.move()
        elif self.mode == 1:
            self.explode()
       

def fireworks_show():
    """ Makes a fireworks show! """
    def step():
        """ Calls up update method on all fireworks in the box 
        and removes if necessary.
        """
        for f in firework_box[:]:
            f.update()
            if f.should_terminate():
                gw.remove(f.get_object())
                firework_box.remove(f)


    def give_me_more_fireworks():
        """ Adds more fireworks to the box. """
        new = Firework(2)
        firework_box.append(new)
        gw.add(new.get_object())

    gw = GWindow(GW_WIDTH, GW_HEIGHT)
    sky = GRect(GW_WIDTH, GW_HEIGHT)
    sky.set_filled(True)
    gw.add(sky)
    firework_box = []

    gw.set_interval(step, 20)
    gw.set_interval(give_me_more_fireworks, 100)

if __name__ == '__main__':
    fireworks_show()