# CelerisLab/common/streakline/_render.py
"""
Render streakline particle clouds to density images.
"""

from __future__ import annotations

import math
import os
from typing import Optional, Sequence, Tuple

import numpy as np

from ._config import CylinderSpec


def gaussian_blur2d(img: np.ndarray, sigma: float = 1.0) -> np.ndarray:
    """Separable Gaussian blur."""
    if float(sigma) <= 0.0:
        return np.asarray(img, dtype=np.float64)
    radius = max(1, int(math.ceil(3.0 * float(sigma))))
    x = np.arange(-radius, radius + 1, dtype=np.float64)
    k = np.exp(-(x**2) / (2.0 * sigma * sigma))
    k /= np.sum(k)
    arr = np.asarray(img, dtype=np.float64)
    for i in range(arr.shape[0]):
        arr[i, :] = np.convolve(arr[i, :], k, mode="same")
    for j in range(arr.shape[1]):
        arr[:, j] = np.convolve(arr[:, j], k, mode="same")
    return arr


def _draw_cylinders(rgb: np.ndarray, cylinders: Sequence[CylinderSpec]) -> None:
    ny, nx = rgb.shape[0], rgb.shape[1]
    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


def render_density(
    positions: np.ndarray,
    ages: np.ndarray,
    *,
    nx: int,
    ny: int,
    out_path: str,
    cylinders: Sequence[CylinderSpec] = (),
    age_decay_steps: float = 20000.0,
    blur_sigma: float = 1.2,
    background_color: tuple[float, float, float] = (0.0, 0.0, 0.0),
    streak_color: tuple[float, float, float] = (1.0, 0.5, 0.0),
    show_release_points: Optional[np.ndarray] = None,
) -> dict:
    """Render streakline particle cloud as a density image.

    Particles are binned onto the pixel grid with age-weighted exponential
    decay, Gaussian-blurred, then composited onto *background_color* using
    *streak_color*.
    """
    density = _build_density_map(
        positions, ages, nx, ny, age_decay_steps=age_decay_steps
    )
    density = gaussian_blur2d(density, sigma=float(blur_sigma))
    return _save_minimal_image(
        density,
        nx,
        ny,
        out_path,
        cylinders,
        background_color=background_color,
        streak_color=streak_color,
        release_points=show_release_points,
    )


def _build_density_map(
    positions: np.ndarray,
    ages: np.ndarray,
    nx: int,
    ny: int,
    *,
    age_decay_steps: float,
) -> np.ndarray:
    if positions.size == 0:
        return np.zeros((ny, nx), dtype=np.float64)
    w = np.exp(-np.asarray(ages, dtype=np.float64) / max(1.0, float(age_decay_steps)))
    y_idx = np.clip(np.floor(positions[:, 1]).astype(np.int64), 0, ny - 1)
    x_idx = np.clip(np.floor(positions[:, 0]).astype(np.int64), 0, nx - 1)
    density = np.zeros((ny, nx), dtype=np.float64)
    np.add.at(density, (y_idx, x_idx), w)
    return density


def _save_minimal_image(
    density: np.ndarray,
    nx: int,
    ny: int,
    out_path: str,
    cylinders: Sequence[CylinderSpec],
    *,
    background_color: tuple[float, float, float],
    streak_color: tuple[float, float, float],
    release_points: Optional[np.ndarray] = None,
) -> dict:
    """Compose density into an RGB image and save (no matplotlib axes)."""
    rgb = np.ones((ny, nx, 3), dtype=np.float64)
    rgb[:, :, 0] = background_color[0]
    rgb[:, :, 1] = background_color[1]
    rgb[:, :, 2] = background_color[2]

    if np.any(density > 0):
        vmax = float(np.percentile(density[density > 0], 99.0))
        alpha = np.clip(density / max(vmax, 1e-12), 0.0, 1.0)
        for c, val in enumerate(streak_color):
            rgb[:, :, c] = np.clip(
                rgb[:, :, c] * (1.0 - alpha) + val * alpha, 0.0, 1.0
            )

    if cylinders:
        _draw_cylinders(rgb, cylinders)

    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,
        "n_particles": int(density.sum()),
        "density_max": float(np.max(density)) if density.size else 0.0,
    }
