# CelerisLab/common/streakline/_integrate.py
"""
Particle integration: bilinear interpolation, RK4 substep, validity filtering.
"""

from __future__ import annotations

import math
from typing import Optional, Sequence

import numpy as np

from ._config import CylinderSpec, FlowFrame, IntegratorConfig


def bilinear_sample(field: np.ndarray, x: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Bilinear interpolation at floating-point positions ``(x, y)``.

    Returns ``NaN`` for positions outside the domain.
    """
    ny, nx = field.shape
    x0 = np.floor(x).astype(np.int64)
    y0 = np.floor(y).astype(np.int64)
    x1 = x0 + 1
    y1 = y0 + 1
    inside = (x0 >= 0) & (x1 < nx) & (y0 >= 0) & (y1 < ny)
    out = np.full_like(x, np.nan, dtype=np.float64)
    if not np.any(inside):
        return out

    xi = x[inside] - x0[inside]
    yi = y[inside] - y0[inside]
    f00 = field[y0[inside], x0[inside]]
    f10 = field[y0[inside], x1[inside]]
    f01 = field[y1[inside], x0[inside]]
    f11 = field[y1[inside], x1[inside]]
    out[inside] = (
        (1.0 - xi) * (1.0 - yi) * f00
        + xi * (1.0 - yi) * f10
        + (1.0 - xi) * yi * f01
        + xi * yi * f11
    )
    return out


def velocity_at(
    pos: np.ndarray, theta: float, f0: FlowFrame, f1: FlowFrame
) -> np.ndarray:
    """Space-time interpolated velocity at positions ``pos`` (N, 2)."""
    x = pos[:, 0]
    y = pos[:, 1]
    vx0 = bilinear_sample(f0.ux, x, y)
    vy0 = bilinear_sample(f0.uy, x, y)
    vx1 = bilinear_sample(f1.ux, x, y)
    vy1 = bilinear_sample(f1.uy, x, y)
    vx = (1.0 - theta) * vx0 + theta * vx1
    vy = (1.0 - theta) * vy0 + theta * vy1
    return np.stack([vx, vy], axis=1)


def particles_valid_mask(
    p_new: np.ndarray,
    *,
    nx: int,
    ny: int,
    cylinders: Sequence[CylinderSpec],
) -> np.ndarray:
    """Boolean mask: particles within domain bounds and outside solids."""
    valid = np.all(np.isfinite(p_new), axis=1)
    valid &= (p_new[:, 0] >= 0.0) & (p_new[:, 0] <= nx - 1)
    valid &= (p_new[:, 1] >= 0.0) & (p_new[:, 1] <= ny - 1)

    for (cx, cy), radius in cylinders:
        dx = p_new[:, 0] - float(cx)
        dy = p_new[:, 1] - float(cy)
        valid &= (dx * dx + dy * dy) > (float(radius) * float(radius))
    return valid


def inject_particles(
    particles: np.ndarray,
    ages: np.ndarray,
    release_points: np.ndarray,
    *,
    inject_per_seed: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Append new particles at *release_points* with age 0."""
    new_pts = np.repeat(release_points, max(1, int(inject_per_seed)), axis=0)
    if particles.size == 0:
        return new_pts.copy(), np.zeros((new_pts.shape[0],), dtype=np.float64)
    return (
        np.vstack([particles, new_pts]),
        np.concatenate([ages, np.zeros((new_pts.shape[0],), dtype=np.float64)]),
    )


def advance_between_frames(
    particles: np.ndarray,
    ages: np.ndarray,
    frame0: FlowFrame,
    frame1: FlowFrame,
    *,
    nx: int,
    ny: int,
    cfg: IntegratorConfig,
    cylinders: Sequence[CylinderSpec],
    rng: Optional[np.random.Generator] = None,
) -> tuple[np.ndarray, np.ndarray, dict, np.ndarray]:
    """Advance particles from *frame0* to *frame1* with RK4.

    Returns:
        particles, ages, diagnostics, survived_mask (input-length bool).
    """
    if particles.size == 0:
        return (
            particles,
            ages,
            {"n_substeps": 0, "max_speed": 0.0},
            np.zeros((0,), dtype=bool),
        )

    dt_save = float(frame1.step - frame0.step)
    max_speed = float(
        max(
            np.nanmax(np.sqrt(frame0.ux * frame0.ux + frame0.uy * frame0.uy)),
            np.nanmax(np.sqrt(frame1.ux * frame1.ux + frame1.uy * frame1.uy)),
        )
    )
    umax = max(max_speed, 1e-8)
    dt_trace = min(float(cfg.alpha_t) * dt_save, float(cfg.alpha_x) / umax)
    n_sub = max(1, int(math.ceil(dt_save / dt_trace)))
    dt_sub = dt_save / float(n_sub)

    p = np.asarray(particles, dtype=np.float64)
    a = np.asarray(ages, dtype=np.float64)
    track = np.ones((p.shape[0],), dtype=bool)
    local_rng = rng or np.random.default_rng(int(cfg.random_seed))
    diffusion = float(cfg.diffusion_coeff)

    for sub_idx in range(n_sub):
        if p.size == 0:
            break
        theta = float(sub_idx) / float(n_sub)
        k1 = velocity_at(p, theta, frame0, frame1)
        k2 = velocity_at(
            p + 0.5 * dt_sub * k1, min(1.0, theta + 0.5 / n_sub), frame0, frame1
        )
        k3 = velocity_at(
            p + 0.5 * dt_sub * k2, min(1.0, theta + 0.5 / n_sub), frame0, frame1
        )
        k4 = velocity_at(
            p + dt_sub * k3, min(1.0, theta + 1.0 / n_sub), frame0, frame1
        )
        p_new = p + (k1 + 2.0 * k2 + 2.0 * k3 + k4) * (dt_sub / 6.0)

        if diffusion > 0.0:
            sigma = math.sqrt(max(0.0, 2.0 * diffusion * dt_sub))
            p_new += local_rng.normal(0.0, sigma, size=p_new.shape)

        valid = particles_valid_mask(p_new, nx=nx, ny=ny, cylinders=cylinders)
        track[track] = valid
        p = p_new[valid]
        a = a[valid] + dt_sub
        if cfg.max_particle_age is not None:
            age_keep = a <= float(cfg.max_particle_age)
            track[track] = age_keep
            p = p[age_keep]
            a = a[age_keep]

    survived_mask = track
    return (
        p,
        a,
        {
            "dt_trace": float(dt_sub),
            "n_substeps": int(n_sub),
            "max_speed": float(max_speed),
        },
        survived_mask,
    )
