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(255,32,0)

@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)
    def scale(self, n: Union[int,float]):
        return Vec2(x=self.x*n, y=self.y*n)

@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.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_speed: float
    particle: list[Particle]
    current_phi: float = 0 # in degree
    max_particle: int = 100
    current_particle: int = 0
    life: int = 100
    step: float = 10
    min_speed_factor: float = 0.5

    def add_particle(self):
        if self.current_particle < self.max_particle:
            new_particle = Particle(
                pos=self.pos,
                v=Vec2(
                    x=math.cos(math.radians(self.current_phi)),
                    y=math.sin(math.radians(self.current_phi)),
                ).scale(self.init_speed),
                life=self.life,
                lifespan=self.life,
            )
            self.particle.append(new_particle)
            self.current_phi = (self.current_phi + self.step) % 360
    
    def update_particle(self):
        for i, p in enumerate(self.particle):
            p.life -= 1
            if p.life <= 0: continue
            p.pos = p.pos + p.v.scale(max(p.life/p.lifespan, self.min_speed_factor))
        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_speed=1,
    particle=[]
)

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

lbl_particle_emit_angle_step = ttk.Label(control_panel, text='Particle emit angle step (deg)')
lbl_particle_emit_angle_step.pack()
def update_emit_width(val):
    lbl_particle_emit_angle_step['text'] = f'Particle emit angle step (deg): {val}'
    em.step = float(val)
sc_particle_emit_angle_step = ttk.Scale(control_panel, orient=tkinter.HORIZONTAL, length=400, from_=0, to=360, command=update_emit_width)
sc_particle_emit_angle_step.pack()
sc_particle_emit_angle_step.set(116)

lbl_particle_emit_speed = ttk.Label(control_panel, text='Particle emit speed')
lbl_particle_emit_speed.pack()
def update_emit_speed_variant(val):
    lbl_particle_emit_speed['text'] = f'Particle emit speed: {val}'
    em.init_speed = float(val)
sc_particle_emit_speed = ttk.Scale(control_panel, orient=tkinter.HORIZONTAL, length=400, from_=-10, to=10, command=update_emit_speed_variant)
sc_particle_emit_speed.pack()
sc_particle_emit_speed.set(3.5)


lbl_particle_emit_min_speed_factor = ttk.Label(control_panel, text='Particle emit min speed factor')
lbl_particle_emit_min_speed_factor.pack()
def update_emit_g(val):
    lbl_particle_emit_min_speed_factor['text'] = f'Particle emit min speed factor: {val}'
    em.min_speed_factor = float(val)
sc_particle_emit_min_speed_factor = ttk.Scale(control_panel, orient=tkinter.HORIZONTAL, length=400, from_=-1, to=1, command=update_emit_g)
sc_particle_emit_min_speed_factor.pack()
sc_particle_emit_min_speed_factor.set(0.5)


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'
)

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

root.mainloop()
