# CelerisLab/common/pathline.py
"""
Particle pathline history accumulation and rendering (debug tool).

This is **not** a streakline — it accumulates each particle's full trajectory
over time.  For streakline visualization use ``streakline.Streakline``.
"""

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from typing import List, Optional, Sequence, Tuple

import numpy as np

from .streakline._render import gaussian_blur2d
from ._types import CylinderSpec


def configure_compute_threads(num_threads: int = 0) -> int:
    """Set BLAS/OpenMP thread env vars for host-side numpy work."""
    import os

    if num_threads <= 0:
        num_threads = min(32, os.cpu_count() or 4)
    for var in (
        "OMP_NUM_THREADS",
        "OPENBLAS_NUM_THREADS",
        "MKL_NUM_THREADS",
        "NUMEXPR_NUM_THREADS",
        "VECLIB_MAXIMUM_THREADS",
    ):
        os.environ[var] = str(int(num_threads))
    return int(num_threads)


@dataclass
class ParticleTrailSet:
    """Growing polylines per particle (pathline history, not streakline)."""

    trails: List[np.ndarray] = field(default_factory=list)
    active: List[bool] = field(default_factory=list)

    def inject(self, points: np.ndarray) -> None:
        pts = np.asarray(points, dtype=np.float64)
        if pts.ndim != 2 or pts.shape[1] != 2:
            raise ValueError("points must be [N,2].")
        for p in pts:
            self.trails.append(np.asarray([[p[0], p[1]]], dtype=np.float64))
            self.active.append(True)

    def keep_indices(self, survived_mask: np.ndarray) -> None:
        mask = np.asarray(survived_mask, dtype=bool).ravel()
        j = 0
        for i, is_active in enumerate(self.active):
            if not is_active:
                continue
            if j >= mask.size:
                break
            if not bool(mask[j]):
                self.active[i] = False
            j += 1

    def append_positions(self, particles: np.ndarray) -> None:
        pts = np.asarray(particles, dtype=np.float64)
        j = 0
        for i, is_active in enumerate(self.active):
            if not is_active:
                continue
            if j >= pts.shape[0]:
                break
            self.trails[i] = np.vstack([self.trails[i], pts[j]])
            j += 1


def _rasterize_segment(
    intensity: np.ndarray,
    x0: float,
    y0: float,
    x1: float,
    y1: float,
    weight: float,
) -> None:
    """Stamp one segment into a scalar intensity map."""
    ny, nx = intensity.shape
    n = int(max(abs(x1 - x0), abs(y1 - y0), 1.0)) + 1
    xs = np.linspace(x0, x1, n)
    ys = np.linspace(y0, y1, n)
    xi = np.clip(np.rint(xs).astype(np.int64), 0, nx - 1)
    yi = np.clip(np.rint(ys).astype(np.int64), 0, ny - 1)
    w = float(max(0.0, min(1.0, weight)))
    np.add.at(intensity, (yi, xi), w)


def _draw_trail(
    intensity: np.ndarray,
    trail: np.ndarray,
    *,
    fade_along_trail: bool,
) -> None:
    if trail.shape[0] < 2:
        return
    nseg = trail.shape[0] - 1
    for i in range(nseg):
        if fade_along_trail:
            t = float(i + 1) / float(nseg)
            weight = 0.15 + 0.85 * t
        else:
            weight = 1.0
        _rasterize_segment(
            intensity,
            float(trail[i, 0]),
            float(trail[i, 1]),
            float(trail[i + 1, 0]),
            float(trail[i + 1, 1]),
            weight,
        )


def render_trails(
    trail_set: ParticleTrailSet,
    *,
    nx: int,
    ny: int,
    out_path: str,
    cylinders: Optional[Sequence[CylinderSpec]] = None,
    blur_sigma: float = 0.0,
    fade_along_trail: bool = True,
    num_threads: int = 0,
    background_color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
    streak_color: Tuple[float, float, float] = (1.0, 0.0, 0.0),
) -> dict:
    """Render accumulated path trails (pathline collage — debug only).

    This is NOT a streakline renderer.  See ``streakline.Streakline``.
    """
    n_threads = configure_compute_threads(num_threads)
    intensity = np.zeros((ny, nx), dtype=np.float64)

    trails = trail_set.trails
    if not trails:
        density_max = 0.0
    else:
        chunk = max(1, (len(trails) + n_threads - 1) // n_threads)

        def _worker(batch: List[np.ndarray]) -> np.ndarray:
            local = np.zeros((ny, nx), dtype=np.float64)
            for tr in batch:
                _draw_trail(local, tr, fade_along_trail=fade_along_trail)
            return local

        batches = [trails[i : i + chunk] for i in range(0, len(trails), chunk)]
        if len(batches) == 1:
            intensity = _worker(batches[0])
        else:
            with ThreadPoolExecutor(max_workers=n_threads) as pool:
                parts = list(pool.map(_worker, batches))
            for part in parts:
                intensity = np.maximum(intensity, part)

        if np.any(intensity > 0):
            vmax = float(np.percentile(intensity[intensity > 0], 99.0))
            intensity = np.clip(intensity / max(vmax, 1e-12), 0.0, 1.0)
        density_max = float(np.max(intensity))

    rgb = np.ones((ny, nx, 3), dtype=np.float64)
    rgb[:, :, 0] = background_color[0]
    rgb[:, :, 1] = background_color[1] * (1.0 - intensity) + streak_color[1] * intensity
    rgb[:, :, 2] = background_color[2] * (1.0 - intensity) + streak_color[2] * intensity
    if streak_color[0] < 0.999:
        rgb[:, :, 0] = background_color[0] * (1.0 - intensity) + streak_color[0] * intensity

    if cylinders:
        yy, xx = np.mgrid[0:ny, 0:nx]
        for (cx, cy), radius in cylinders:
            mask = (xx - float(cx)) ** 2 + (yy - float(cy)) ** 2 <= float(radius) ** 2
            rgb[mask] = 0.0

    if float(blur_sigma) > 0.0:
        for c in range(3):
            rgb[:, :, c] = gaussian_blur2d(rgb[:, :, c], sigma=float(blur_sigma))

    rgb = np.clip(rgb, 0.0, 1.0)
    import os
    os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)

    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        fig = plt.figure(figsize=(nx / 100.0, ny / 100.0), dpi=100)
        ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])
        ax.axis("off")
        ax.imshow(rgb, origin="lower", interpolation="nearest", aspect="auto")
        fig.savefig(out_path, dpi=100, bbox_inches="tight", pad_inches=0)
        plt.close(fig)
    except ImportError:
        from PIL import Image
        Image.fromarray((np.clip(rgb, 0.0, 1.0) * 255.0).astype(np.uint8)).save(out_path)

    return {
        "image_path": out_path,
        "trail_count": int(len(trails)),
        "density_max": density_max,
        "num_threads": int(n_threads),
    }
