import math
import tkinter
from tkinter import ttk
from dataclasses import dataclass
from random import random
from typing import Union
import time

WINDOW_WIDTH = 300
WINDOW_HEIGHT = 300

root = tkinter.Tk()
root.title('Simple particle system')
canvas_main = tkinter.Canvas(root, bg="black", height=WINDOW_HEIGHT, width=WINDOW_WIDTH)

@dataclass
class Color:
    r: int
    g: int
    b: int
    a: int = 0xff  # 0xff-solid 0x00-transparent

    def get_rgb_string(self) -> str:
        return f'#{self.r:02x}{self.g:02x}{self.b:02x}'

    def get_rgba_string(self) -> str:
        return f'#{self.r:02x}{self.g:02x}{self.b:02x}{self.a:02x}'

    def scale(self, light: float):
        return Color(
            r=math.floor(self.r * light),
            g=math.floor(self.g * light),
            b=math.floor(self.b * light),
            a=self.a,
        )

DEFAULT_PARTICLE_COLOR = Color(0,255,255)

@dataclass
class Vec2:
    x: Union[int, float]
    y: Union[int, float]
    def __add__(self, other):
        return Vec2(x=self.x+other.x, y=self.y+other.y)
    def length(self) -> float:
        return math.sqrt(self.x**2+self.y**2)

@dataclass
class NumVariant:
    base: Union[int, float]
    variant: Union[int, float]

    def value(self):
        return (float if type(self.base) is float else int)(self.base + (self.variant * random()))

PARTICLE_SIZE: int = 4

@dataclass
class Particle:
    pos: Vec2
    v: Vec2
    life: int
    lifespan: int
    color: Color = DEFAULT_PARTICLE_COLOR

    def __post_init__(self):
        color_str = self.color.get_rgb_string()
        self._elem = canvas_main.create_oval(
            self.pos.x-PARTICLE_SIZE/2, self.pos.y-PARTICLE_SIZE/2,
            self.pos.x+PARTICLE_SIZE/2, self.pos.y+PARTICLE_SIZE/2,
            fill=color_str,
            outline=color_str,
        )

    def draw(self):
        canvas_main.moveto(self._elem,
            self.pos.x-PARTICLE_SIZE/2, self.pos.y-PARTICLE_SIZE/2,
        )
        color_str = self.color.scale(self.life/self.lifespan).get_rgb_string()
        canvas_main.itemconfigure(self._elem,
            fill=color_str,
            outline=color_str,
        )

@dataclass
class Vec2Variant:
    v: Vec2
    angle_variant: float  # in radiant; abs val.
    speed_variant: float
    def value(self):
        z = random()*(self.angle_variant*2)-self.angle_variant
        va = math.atan2(self.v.y, self.v.x)
        na = z + va
        s = self.v.length()
        sv = random()*(self.speed_variant*2)-self.speed_variant
        ns = s + sv
        return Vec2(x=math.cos(na)*ns, y=math.sin(na)*ns)

@dataclass
class Emitter:
    name: str
    pos: Vec2
    init_v: Vec2Variant
    particle: list[Particle]
    max_particle: int = 100
    current_particle: int = 0
    life: NumVariant = NumVariant(base=100, variant=2)
    force: Vec2 = Vec2(0, -0.005)

    def add_particle(self):
        if self.current_particle < self.max_particle:
            life = self.life.value()
            new_particle = Particle(
                pos=self.pos,
                v=self.init_v.value(),
                life=life,
                lifespan=life,
            )
            self.particle.append(new_particle)
    
    def update_particle(self):
        for i, p in enumerate(self.particle):
            p.life -= 1
            if p.life <= 0: continue
            p.v = p.v + self.force
            p.pos = p.pos + p.v
        for p in self.particle:
            if (p.life <= 0):
                canvas_main.delete(p._elem)
        self.particle = [p for p in self.particle if p.life > 0]

    def draw_particle(self):
        for p in self.particle: p.draw()

em = Emitter(
    name='em1',
    pos=Vec2(WINDOW_WIDTH/2, WINDOW_HEIGHT/2),
    init_v=Vec2Variant(v=Vec2(0, -0.33), angle_variant=math.radians(45), speed_variant=0),
    particle=[]
)

canvas_main.pack()
canvas_main.create_oval(
    WINDOW_WIDTH/2-PARTICLE_SIZE/2, WINDOW_HEIGHT/2-PARTICLE_SIZE/2,
    WINDOW_WIDTH/2+PARTICLE_SIZE/2, WINDOW_HEIGHT/2+PARTICLE_SIZE/2,
    fill='red', outline='white'
)

control_panel = tkinter.Toplevel()
control_panel.title('Control Panel')

lbl_particle_emit_width = ttk.Label(control_panel, text='Particle emit angle')
lbl_particle_emit_width.pack()
def update_emit_width(val):
    lbl_particle_emit_width['text'] = f'Particle emit angle: {val}'
    em.init_v.angle_variant = math.radians(float(val))
sc_particle_emit_width = ttk.Scale(control_panel, orient=tkinter.HORIZONTAL, length=400, from_=0, to=90, command=update_emit_width)
sc_particle_emit_width.pack()
sc_particle_emit_width.set(30)

lbl_particle_emit_speed_variant = ttk.Label(control_panel, text='Particle emit speed variant')
lbl_particle_emit_speed_variant.pack()
def update_emit_speed_variant(val):
    lbl_particle_emit_speed_variant['text'] = f'Particle emit speed variant: {val}'
    em.init_v.speed_variant = float(val)
sc_particle_emit_speed_variant = ttk.Scale(control_panel, orient=tkinter.HORIZONTAL, length=400, from_=-5, to=5, command=update_emit_speed_variant)
sc_particle_emit_speed_variant.pack()
sc_particle_emit_speed_variant.set(0.5)

lbl_particle_emit_v = ttk.Label(control_panel, text='Particle emit velocity (on y axis)')
lbl_particle_emit_v.pack()
def update_emit_v(val):
    lbl_particle_emit_v['text'] = f'Particle emit velocity (on y axis): {val}'
    em.init_v.v.y = float(val)
sc_particle_emit_v = ttk.Scale(control_panel, orient=tkinter.HORIZONTAL, length=400, from_=-5, to=5, command=update_emit_v)
sc_particle_emit_v.pack()
sc_particle_emit_v.set(-1)

lbl_particle_emit_g = ttk.Label(control_panel, text='Gravity (on y axis)')
lbl_particle_emit_g.pack()
def update_emit_g(val):
    lbl_particle_emit_g['text'] = f'Gravity (on y axis): {val}'
    em.force.y = float(val)
sc_particle_emit_g = ttk.Scale(control_panel, orient=tkinter.HORIZONTAL, length=400, from_=-1, to=1, command=update_emit_g)
sc_particle_emit_g.pack()
sc_particle_emit_g.set(0.05)

while True:
    em.add_particle()
    em.update_particle()
    em.draw_particle()
    root.update()
    time.sleep(1/60)

root.mainloop()