"""Position-Based Fluids in 3D (Macklin & Müller 2013).

Same skeleton as everything else in this project:
    per-particle LOCAL PROJECTION (tied)  ->  GLOBAL RECONCILIATION
The local projection is the density constraint lambda — one rule shared
by every particle (the fluid's tied projector, roadmap §3). The global
step is the position-correction scatter + neighbor grid.

lambda_fn hook: the analytic lambda = -C / (sum|gradC|^2 + eps) can be
replaced by a learned rule (the fluid-token experiment, W7).
"""
import numpy as np
from scipy.spatial import cKDTree


class PBF3D:
    def __init__(self, positions, bounds_min, bounds_max, h=0.1, rho0=1000.0,
                 dt=0.008, iters=4, gravity=(0.0, -9.81, 0.0),
                 scorr_k=1e-4, xsph_c=0.05, eps=100.0, lambda_fn=None):
        self.x = np.asarray(positions, dtype=np.float64).copy()
        self.v = np.zeros_like(self.x)
        self.lo = np.asarray(bounds_min, dtype=np.float64)
        self.hi = np.asarray(bounds_max, dtype=np.float64)
        self.h, self.rho0, self.dt, self.iters = h, rho0, dt, iters
        self.gravity = np.asarray(gravity, dtype=np.float64)
        self.scorr_k, self.xsph_c, self.eps = scorr_k, xsph_c, eps
        self.lambda_fn = lambda_fn

        self.poly6_c = 315.0 / (64.0 * np.pi * h**9)
        self.spiky_c = -45.0 / (np.pi * h**6)
        self.w_dq = self.poly6_c * (h**2 - (0.3 * h)**2) ** 3  # scorr ref

        # particle mass so a filled lattice sits at rest density: the most
        # crowded particle of the initial configuration defines rho0
        S = self._kernel_sums(self.x)
        self.mass = rho0 / S.max()
        self.rho = np.full(len(self.x), rho0)

    def _kernel_sums(self, p):
        tree = cKDTree(p)
        pairs = tree.query_pairs(self.h, output_type="ndarray")
        r2 = ((p[pairs[:, 0]] - p[pairs[:, 1]]) ** 2).sum(axis=1)
        w = self.poly6_c * (self.h**2 - r2) ** 3
        S = np.full(len(p), self.poly6_c * self.h**6)  # self term W(0)
        np.add.at(S, pairs[:, 0], w)
        np.add.at(S, pairs[:, 1], w)
        return S

    def step(self, record=None):
        """One PBF step. record: optional list collecting (C, g2, lam)
        rows from the final constraint iteration (corpus hook)."""
        h, dt, rho0, m = self.h, self.dt, self.rho0, self.mass
        v = self.v + dt * self.gravity
        p = self.x + dt * v
        np.clip(p, self.lo, self.hi, out=p)

        tree = cKDTree(p)
        pairs = tree.query_pairs(h, output_type="ndarray")
        i, j = pairs[:, 0], pairs[:, 1]
        N = len(p)

        for it in range(self.iters):
            d = p[i] - p[j]
            r2 = (d * d).sum(axis=1)
            r = np.sqrt(np.maximum(r2, 1e-12))
            w = self.poly6_c * np.maximum(h**2 - r2, 0.0) ** 3
            gw = (self.spiky_c * np.maximum(h - r, 0.0) ** 2 / r)[:, None] * d
            gw *= m / rho0  # fold mass: gw is now a grad-C contribution

            rho = np.full(N, m * self.poly6_c * h**6)
            np.add.at(rho, i, m * w)
            np.add.at(rho, j, m * w)
            C = rho / rho0 - 1.0

            sumg = np.zeros_like(p)
            np.add.at(sumg, i, gw)
            np.add.at(sumg, j, -gw)
            g2 = (sumg * sumg).sum(axis=1)
            gw2 = (gw * gw).sum(axis=1)
            np.add.at(g2, i, gw2)
            np.add.at(g2, j, gw2)

            if self.lambda_fn is not None:
                lam = self.lambda_fn(C, g2)
            else:
                lam = -C / (g2 + self.eps * (m / rho0) ** 2)

            scorr = -self.scorr_k * (w * (m / rho0) / self.w_dq) ** 4
            coef = (lam[i] + lam[j] + scorr)[:, None] * gw
            dp = np.zeros_like(p)
            np.add.at(dp, i, coef)
            np.add.at(dp, j, -coef)
            p += dp
            np.clip(p, self.lo, self.hi, out=p)

        if record is not None:
            record.append(np.stack([C, g2, lam], axis=1))

        self.v = (p - self.x) / dt
        # XSPH viscosity on the final neighbor set
        dv = self.v[j] - self.v[i]
        wv = (self.poly6_c * np.maximum(
            h**2 - ((p[i] - p[j]) ** 2).sum(axis=1), 0.0) ** 3
            * m / rho0)[:, None] * dv
        xs = np.zeros_like(p)
        np.add.at(xs, i, wv)
        np.add.at(xs, j, -wv)
        self.v += self.xsph_c * xs
        self.x = p
        self.rho = rho
        return C

    def density_error(self):
        """Mean |rho/rho0 - 1| over particles at or above rest density
        (free surfaces are legitimately under-dense in SPH kernels)."""
        c = self.rho / self.rho0 - 1.0
        return np.abs(c[c > -0.05]).mean()


def dam_break_block(lo, extent, spacing, jitter=1e-3, seed=0):
    """Lattice block of particles for dam-break initial conditions."""
    rng = np.random.default_rng(seed)
    axes = [np.arange(lo[k] + spacing / 2, lo[k] + extent[k], spacing)
            for k in range(3)]
    g = np.stack(np.meshgrid(*axes, indexing="ij"), axis=-1).reshape(-1, 3)
    return g + rng.normal(scale=jitter, size=g.shape)