"""Generate a synthetic top-down video of flies walking in a circular arena.

Each fly does a correlated random walk and alternates between active and rest
(sleep-like) states via a 2-state Markov process; it bounces off the arena wall.
Flies are dark Gaussian blobs on a bright background. Self-contained — no external
(licence-encumbered) video needed. Also returns the ground-truth tracks.
"""
from __future__ import annotations

import numpy as np


def fly_arena(T: int = 300, H: int = 256, W: int = 256, n_flies: int = 8,
              fps: float = 15.0, seed: int = 0):
    rng = np.random.default_rng(seed)
    cx, cy, R = W / 2, H / 2, min(H, W) / 2 - 12

    # initial positions inside the arena
    ang = rng.uniform(0, 2 * np.pi, n_flies)
    rad = R * np.sqrt(rng.uniform(0, 0.7, n_flies))
    x = cx + rad * np.cos(ang)
    y = cy + rad * np.sin(ang)
    heading = rng.uniform(0, 2 * np.pi, n_flies)
    active = np.ones(n_flies, bool)

    tracks = np.zeros((T, n_flies, 2), np.float32)   # ground truth (x, y)
    states = np.zeros((T, n_flies), np.uint8)        # 1 = active

    yy, xx = np.mgrid[0:H, 0:W]
    frames = np.empty((T, H, W), np.uint8)

    for t in range(T):
        # Markov active/rest transitions
        flip_to_rest = active & (rng.random(n_flies) < 0.012)
        flip_to_active = (~active) & (rng.random(n_flies) < 0.04)
        active = (active & ~flip_to_rest) | flip_to_active

        heading += rng.normal(0, 0.4, n_flies)
        speed = np.where(active, rng.normal(2.6, 0.5, n_flies).clip(0.5, 5), 0.0)
        x = x + speed * np.cos(heading)
        y = y + speed * np.sin(heading)

        # bounce off the circular wall
        dx, dy = x - cx, y - cy
        dist = np.hypot(dx, dy)
        out = dist > R
        if out.any():
            nrm = np.arctan2(dy[out], dx[out])
            heading[out] = nrm + np.pi + rng.normal(0, 0.3, out.sum())
            x[out] = cx + (R - 1) * np.cos(nrm)
            y[out] = cy + (R - 1) * np.sin(nrm)

        tracks[t, :, 0] = x
        tracks[t, :, 1] = y
        states[t] = active.astype(np.uint8)

        # render: bright background, arena ring, dark fly blobs
        frame = np.full((H, W), 225.0, np.float32)
        ring = np.abs(np.hypot(xx - cx, yy - cy) - R) < 1.5
        frame[ring] = 150.0
        for fx, fy in zip(x, y):
            d2 = (xx - fx) ** 2 + (yy - fy) ** 2
            frame -= 185.0 * np.exp(-d2 / (2 * 3.0 ** 2))
        frame += rng.normal(0, 3, (H, W))
        frames[t] = np.clip(frame, 0, 255).astype(np.uint8)

    meta = {"fps": fps, "arena": [float(cx), float(cy), float(R)], "n_flies": n_flies}
    return frames, tracks, states, meta