"""Synthetic media generators shared by tests, data building and benchmarks."""
import numpy as np
from PIL import Image, ImageDraw

COLORS = {"red": (220, 40, 40), "blue": (40, 80, 220), "green": (40, 170, 60), "yellow": (230, 200, 30),
          "purple": (140, 60, 180), "orange": (240, 130, 30)}
SHAPES = ["circle", "square", "triangle"]


def draw_shape(d, shape, color, cx, cy, r):
    c = COLORS[color]
    if shape == "circle":
        d.ellipse([cx - r, cy - r, cx + r, cy + r], fill=c)
    elif shape == "square":
        d.rectangle([cx - r, cy - r, cx + r, cy + r], fill=c)
    else:
        d.polygon([(cx, cy - r), (cx - r, cy + r), (cx + r, cy + r)], fill=c)


def shapes_image(objs, size=384, bg=(245, 245, 240)):
    """objs: list of (shape, color, cx, cy, r) in [0,1] coords for cx, cy and fraction r."""
    im = Image.new("RGB", (size, size), bg)
    d = ImageDraw.Draw(im)
    for shape, color, cx, cy, r in objs:
        draw_shape(d, shape, color, cx * size, cy * size, r * size)
    return im


def beeps(n, sr=16000, dur=4.0, freq=880.0, rng=None):
    """n short tone beeps spread over `dur` seconds, plus light noise."""
    rng = rng or np.random.default_rng(0)
    t = np.zeros(int(sr * dur), np.float32)
    if n:
        starts = np.linspace(0.3, dur - 0.5, n) + rng.uniform(-0.05, 0.05, n)
        for s in starts:
            a = int(s * sr); L = int(0.18 * sr)
            tt = np.arange(L) / sr
            t[a:a + L] += 0.5 * np.sin(2 * np.pi * freq * tt) * np.hanning(L)
    t += rng.normal(0, 0.01, t.shape).astype(np.float32)
    return t


def moving_video(shape, color, direction, n_frames=8, size=320, flashes=0, change_to=None, rng=None):
    """A shape moving across the frame. flashes: number of frames where the background flashes white-yellow.
    change_to: the colour the shape switches to halfway through (None = no change)."""
    rng = rng or np.random.default_rng(0)
    frames = []
    x0, y0 = rng.uniform(0.35, 0.65), rng.uniform(0.35, 0.65)
    dx = {"left": (-1, 0), "right": (1, 0), "up": (0, -1), "down": (0, 1)}[direction]
    span = 0.6
    flash_idx = set(rng.choice(np.arange(1, n_frames - 1), size=flashes, replace=False).tolist()) if flashes else set()
    for i in range(n_frames):
        f = i / (n_frames - 1) - 0.5
        cx = x0 + dx[0] * f * span; cy = y0 + dx[1] * f * span
        cx = min(max(cx, 0.1), 0.9); cy = min(max(cy, 0.1), 0.9)
        bg = (255, 250, 150) if i in flash_idx else (40, 40, 48)
        col = change_to if (change_to and i >= n_frames // 2) else color
        frames.append(shapes_image([(shape, col, cx, cy, 0.09)], size=size, bg=bg))
    return frames