import math, random, time
from dataclasses import dataclass, field
from typing import List, Optional, Tuple
import gradio as gr
from PIL import Image, ImageDraw

W, H = 900, 540
ROOM = (60, 40, W - 60, H - 40)
DT = 0.05
MAX_SPEED_BASE = 65.0
PERSON_RADIUS = 6
EXIT_W, EXIT_H = 28, 68
C_BG = (245, 246, 248, 255)
C_ROOM = (250, 250, 250, 255)
C_BORDER = (28, 28, 28, 255)
C_OBS = (175, 178, 185, 255)
C_EXIT = (52, 132, 255, 255)
C_CALM = (44, 184, 78, 255)
C_ALERT = (250, 191, 35, 255)
C_PANIC = (233, 68, 68, 255)
C_INCIDENT = (230, 60, 60, 96)

RUNNING = False

@dataclass
class Person:
    x: float; y: float; vx: float; vy: float
    stress: float = 0.0; state: int = 0; evacuated: bool = False
    def color(self): return [C_CALM, C_ALERT, C_PANIC][self.state]

@dataclass
class ExitZone:
    x: int; y: int; w: int = EXIT_W; h: int = EXIT_H
    def rect(self): return (self.x, self.y, self.x + self.w, self.y + self.h)

@dataclass
class World:
    width: int = W; height: int = H; room: Tuple[int,int,int,int] = ROOM
    people: List[Person] = field(default_factory=list)
    exits: List[ExitZone] = field(default_factory=list)
    incident: Optional[Tuple[int,int,int]] = None
    speed_scale: float = 1.0; cohesion: float = 0.3; separation: float = 0.6
    panic_spread: float = 2.0; personal_space: float = 22.0

    def step(self, dt: float):
        def nearest_exit_vec(px, py):
            best=None; best_d2=1e12
            for ex in self.exits:
                cx, cy = ex.x + ex.w/2, ex.y + ex.h/2
                dx, dy = cx - px, cy - py
                d2 = dx*dx + dy*dy
                if d2 < best_d2: best_d2, best = d2, (dx, dy)
            return best or (0.0, 0.0)

        for p in self.people:
            if p.evacuated: continue
            gx, gy = nearest_exit_vec(p.x, p.y)
            p.vx += gx*0.002; p.vy += gy*0.002
            v = math.hypot(p.vx, p.vy)
            vmax = MAX_SPEED_BASE*self.speed_scale
            if v > vmax: p.vx, p.vy = p.vx/v*vmax, p.vy/v*vmax
            p.x += p.vx*dt; p.y += p.vy*dt
            for ex in self.exits:
                x0, y0, x1, y1 = ex.rect()
                if x0 <= p.x <= x1 and y0 <= p.y <= y1:
                    p.evacuated = True

    def render(self):
        img = Image.new("RGBA", (self.width, self.height), C_BG)
        d = ImageDraw.Draw(img, "RGBA")
        d.rectangle(self.room, fill=C_ROOM, outline=C_BORDER, width=3)
        for ex in self.exits:
            d.rectangle(ex.rect(), fill=C_EXIT)
            d.text((ex.x+6, ex.y+ex.h/2-7), "Exit", fill=(255,255,255,255))
        for p in self.people:
            if not p.evacuated:
                d.ellipse((p.x-4, p.y-4, p.x+4, p.y+4), fill=p.color())
        return img

def make_world(n_people=200):
    w = World()
    ex1 = ExitZone(x=ROOM[2]-8-EXIT_W, y=int(ROOM[1]+70))
    ex2 = ExitZone(x=ROOM[2]-8-EXIT_W, y=int(ROOM[3]-70-EXIT_H))
    w.exits = [ex1, ex2]
    for _ in range(n_people):
        w.people.append(Person(random.randint(100, 700), random.randint(100, 400), 0, 0))
    return w

def reset_fn(n_people):
    w = make_world(int(n_people))
    return w.render(), w

def start_fn(world):
    global RUNNING; RUNNING = True
    while RUNNING:
        world.step(DT)
        yield world.render(), world
        time.sleep(DT)

def pause_fn(): 
    global RUNNING; RUNNING = False
    return gr.update(), gr.update()

with gr.Blocks(title="Behavioral ML in Virtual Crowds") as demo:
    gr.Markdown("## Behavioral ML in Virtual Crowds · eenvoudige simulatie")
    img = gr.Image(label="Simulatie")
    n_people = gr.Slider(10, 500, value=200, step=10, label="Aantal mensen")
    start_btn = gr.Button("▶️ Start")
    pause_btn = gr.Button("⏸️ Pauze")
    reset_btn = gr.Button("🔁 Reset")
    state = gr.State()
    reset_btn.click(reset_fn, [n_people], [img, state])
    start_btn.click(start_fn, [state], [img, state])
    pause_btn.click(pause_fn, outputs=[img, state])

if __name__ == "__main__":
    demo.launch()