# CelerisLab/common/streakline/_streakline.py
"""
Streakline — passive-consumer particle-based flow visualization.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional, Sequence

import numpy as np

from ._config import (
    CylinderSpec,
    FlowFrame,
    IntegratorConfig,
    ReleaseConfig,
    build_release_points,
)
from ._integrate import advance_between_frames, inject_particles
from ._render import render_density


@dataclass
class Streakline:
    """Passive streakline integrator.

    Does **not** own a simulation or control the stepping loop.  Users call
    ``observe()`` whenever they have a new velocity frame; body actions,
    checkpointing, and other modules are managed independently.

    Usage::

        streak = Streakline(release_points=..., nx=nx, ny=ny, ...)
        # In your own simulation loop:
        macro = sim.get_macroscopic()
        streak.observe(ux=macro["ux"], uy=macro["uy"], step=step)
        streak.render("output.png")
    """

    release_points: np.ndarray
    release_cfg: ReleaseConfig
    integrator_cfg: IntegratorConfig
    nx: int
    ny: int
    cylinders: Sequence[CylinderSpec] = field(default_factory=list)

    _particles: np.ndarray = field(init=False, repr=False)
    _ages: np.ndarray = field(init=False, repr=False)
    _prev_frame: Optional[FlowFrame] = field(init=False, repr=False, default=None)

    def __post_init__(self):
        self._particles = np.zeros((0, 2), dtype=np.float64)
        self._ages = np.zeros((0,), dtype=np.float64)
        self._dense_pts = build_release_points(
            np.asarray(self.release_points), self.release_cfg
        )

    def observe(self, ux: np.ndarray, uy: np.ndarray, step: int) -> dict:
        """Feed one velocity snapshot.

        First call caches the frame.  From the second call onward particles
        are injected and advanced from the previous frame to this one.
        """
        frame = FlowFrame(
            step=int(step),
            ux=np.asarray(ux, dtype=np.float64),
            uy=np.asarray(uy, dtype=np.float64),
        )

        if self._prev_frame is None:
            self._prev_frame = frame
            return {}

        self._particles, self._ages = inject_particles(
            self._particles,
            self._ages,
            self._dense_pts,
            inject_per_seed=self.release_cfg.inject_per_seed,
        )

        self._particles, self._ages, diag, _ = advance_between_frames(
            self._particles,
            self._ages,
            self._prev_frame,
            frame,
            nx=self.nx,
            ny=self.ny,
            cfg=self.integrator_cfg,
            cylinders=self.cylinders,
        )

        self._prev_frame = frame
        return diag

    def render(self, out_path: str, **kwargs) -> dict:
        """Render current particle cloud as a density image."""
        # Default age_decay_steps from integrator config; **kwargs may override
        kw = dict(kwargs, cylinders=self.cylinders)
        kw.setdefault("age_decay_steps", self.integrator_cfg.age_decay_steps)
        return render_density(
            self._particles,
            self._ages,
            nx=self.nx,
            ny=self.ny,
            out_path=out_path,
            **kw,
        )

    @property
    def n_particles(self) -> int:
        return self._particles.shape[0]

    def reset(self):
        self._particles = np.zeros((0, 2), dtype=np.float64)
        self._ages = np.zeros((0,), dtype=np.float64)
        self._prev_frame = None
