"""Fireworks.

This module implements a fireworks that is drawn onto a tkinter canvas.

References:

.. _Lecture:
   http://www.informatik.uni-freiburg.de/~ki/teaching/ws1617/info1/lecture.html

Authors:
    Tim Schulte
    Thorsten Engesser

Version:
    WS 2016/17

"""

import tkinter as tk
from time import time, sleep
from random import choice, uniform
from math import sin, cos
from sys import modules

GRAVITY = 30  # you can play around with this if you want


class Particle:
    """Generic class for particles.

    Particles can be emitted by Fireworks objects. They are displayed for a
    specified lifespan and then removed from the canvas.

    Attributes:
        cv (Tk.canvas): the canvas in which the particle is drawn.
        cid (Tk.canvas): the tkinter canvas id for the particle.
        x (float): x-coordinate of the particle.
        y (float): y-coordinate of the particle.
        vx (float): x-velocity of the particle (in pixels per second).
        vy (float): y-velocity of the particle (in pixels per second).
        color (str): color of the particle.
        age (float): age of the particle.
        lifespan (float): lifespan of the particle (in seconds).

    """

    def __init__(self, cv=None, color='white', x=0., y=0.,
                 vx=0., vy=0., lifespan=5.):
        """Init Particle objects.

        Args:
            cv (Tk.canvas): the canvas in which the particle is drawn.
            x (float): x-coordinate of the particle.
                Defaults to 0.0.
            y (float): y-coordinate of the particle.
                Defaults to 0.0.
            vx (float): x-velocity of the particle (in pixels per second).
                Defaults to 0.0.
            vy (float): y-velocity of the particle (in pixels per second).
                Defaults to 0.0.
            color (str): color of the particle.
                Defaults to 'white'.
            lifespan (float): lifespan of the particle (in seconds).
                Defaults to 5.0.

        """
        self.cv = cv
        self.cid = None
        self.x, self.y = x, y
        self.vx, self.vy = vx, vy
        self.color = color
        self.age, self.lifespan = 0, lifespan

    def update(self, dt):
        """Update position and velocity after dt seconds have passed.

        Args:
            dt (float): the time that has passed after the last update (in s).

        """
        self.age += dt
        if self.alive():
            self.vy += GRAVITY * dt
            self.x += self.vx * dt
            self.y += self.vy * dt
            self.cv.move(self.cid, self.vx * dt, self.vy * dt)
        elif self.cid is not None:
            cv.delete(self.cid)
            self.cid = None

    def alive(self):
        """Check if particle is still within its lifespan."""
        return self.age <= self.lifespan


class SquareParticle(Particle):
    """A Particle with a quadratic shape"""
    def __init__(self, x=0., y=0., size=2., **kwargs):
        super().__init__(x=x, y=y, **kwargs)
        self.cid = self.cv.create_polygon(
            x - size, y - size, x + size, y - size,
            x + size, y + size, x - size, y + size,
            fill=self.color)


class TriangleParticle(Particle):
    """A Particle with a triangular shape"""
    def __init__(self, x=0., y=0., size=2., **kwargs):
        super().__init__(x=x, y=y, **kwargs)
        self.cid = self.cv.create_polygon(
            x - size, y - size, x + size,
            y - size, x, y + size,
            fill=self.color)


class CircularParticle(Particle):
    """A Particle with a circular shape."""
    def __init__(self, x=0., y=0., size=2., **kwargs):
        super().__init__(x=x, y=y, **kwargs)
        self.cid = self.cv.create_oval(
            x - size, y - size, x + size,
            y + size, fill=self.color)


class Fireworks:
    """Generic class for fireworks.

    The main "behavior" of a fireworks is specified via its update method.
    E.g., new particles can be emitted and added to the particle list. The
    Fireworks base class automatically updates all particles from the particle
    list in its update method.

    Attributes:
        cv (Tk.canvas): the canvas in which the fireworks is drawn.
        age (float): age of the fireworks.
        particles (list of Particle): list of generated particles.

    """

    def __init__(self, cv=None):
        """Init Fireworks objects.

        Args:
            cv (Tk.canvas): the canvas in which the particle is drawn.

        """
        self.cv = cv
        self.age = 0
        self.particles = []

    def update(self, dt):
        """Update the fireworks' particles and remove dead ones.

        Args:
            dt (float): the time that has passed after the last update (in s).

        """
        self.age += dt
        for p in self.particles:
            p.update(dt)
        for i in range(len(self.particles) - 1, -1, -1):
            if not self.particles[i].alive():
                del self.particles[i]


class Volcano(Fireworks):
    """A volcano that continuously emits colored particles.

    Attributes:
        x (float): x-coordinate of the volcano.
        pps (float): the number of particles to spawn per second.
        colors (list of string): the colors of the particles to spawn.

    """

    def __init__(self, cv, x, pps, colors):
        """Init Volcano objects.

        Args:
            cv (Tk.canvas): the canvas in which the particle is drawn.
            x (float): x-coordinate of the volcano.
            pps (float): the number of particles to spawn per second.
            colors (list of string): the colors of the particles to spawn.

        """
        super().__init__(cv)
        self.cid = cv.create_polygon(x - 12, 530,  # size and color are fixed
                                     x + 12, 530,  # (can be parametrized)
                                     x, 500,
                                     fill="orange")
        self.x = x
        self.pps = pps
        self.colors = colors
        self._tospawn = 0

    def update(self, dt):
        """Continuously emits new random particles and updates them.

        Args:
            dt (float): the time that has passed after the last update (in s).

        """
        super().update(dt)
        self._tospawn += self.pps * dt
        color = self.colors[int(self.age / 3) % len(self.colors)]
        for i in range(int(self._tospawn)):
            ptype = choice(
                [SquareParticle, TriangleParticle, CircularParticle])
            angle = uniform(-0.25, 0.25)
            speed = -uniform(80.0, 120.0)
            vx = sin(angle) * speed
            vy = cos(angle) * speed
            self.particles.append(
                ptype(cv=self.cv, x=self.x, y=500, color=color, vx=vx, vy=vy))
        self._tospawn -= int(self._tospawn)


def simulate(cv, objects):
    """Fireworks simulation loop.

        Args:
            cv (float): the canvas in which the firework objects are drawn.
            objects (float): the firework objects.

        """
    t = time()
    while running:
        sleep(0.01)
        tnew = time()
        t, dt = tnew, tnew - t
        for o in objects:
            o.update(dt)
        cv.update()


def close(*ignore):
    """Stops simulation loop and closes the window."""
    global running
    running = False
    root.destroy()


if __name__ == '__main__':
    root = tk.Tk()
    cv = tk.Canvas(root, height=600, width=800)
    cv.create_rectangle(0, 0, 800, 600, fill="midnight blue")  # sky
    cv.create_rectangle(0, 450, 800, 600, fill="gray11")  # ground
    cv.pack()

    v1 = Volcano(cv, 400, 100, ["red", "green", "silver", "gold"])
    objects = [v1]

    # close with [ESC] or (x) button
    root.bind('<Escape>', close)
    root.protocol("WM_DELETE_WINDOW", close)

    running = True
    root.after(500, simulate, cv, objects)
    if "idlelib" not in modules:
        root.mainloop()