"""Real-time probe-induced soft-tissue deformation (first-order model).

Models the tissue displacement produced by inserting a cylindrical probe along
a trajectory using an **incompressible cavity-expansion** approximation: as a
cavity of radius `a` opens, a point initially at radial distance r from the
probe axis is pushed outward to sqrt(r^2 + a^2) (area/volume-preserving), with a
smooth spatial window so the effect stays local. This is a recognised
first-order model of needle/probe insertion mechanics.

IMPORTANT: this is a fast analytical approximation for planning preview and
demonstration — NOT a validated, patient-specific finite-element/AI deformation
model. The validated model (trained on FEM and verified against measured tissue
displacement) is the funded development described in docs/VALIDATION_PLAN.md.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from .planning import Trajectory


@dataclass
class ProbeDeformation:
    trajectory: Trajectory
    probe_radius_mm: float = 1.5        # cavity radius a
    influence_radius_mm: float = 18.0   # beyond this, displacement -> 0

    def _closest_on_axis(self, P):
        E, T = self.trajectory.entry, self.trajectory.target
        v = T - E
        L = float(np.linalg.norm(v))
        vhat = v / L if L > 0 else v
        t = np.clip((P - E) @ vhat, 0.0, L)
        return E[None, :] + t[:, None] * vhat[None, :]

    def field(self, points):
        """Displacement vectors (N,3 mm) for input world points (N,3)."""
        P = np.atleast_2d(np.asarray(points, dtype=np.float64))
        closest = self._closest_on_axis(P)
        radial = P - closest
        d = np.linalg.norm(radial, axis=1)
        a, R = self.probe_radius_mm, self.influence_radius_mm
        # incompressible cavity expansion: u_r = sqrt(d^2 + a^2) - d
        u = np.sqrt(d**2 + a**2) - d
        # smooth localisation window (1 at axis, 0 at influence radius)
        w = np.where(d < R, 0.5 * (1 + np.cos(np.pi * np.clip(d / R, 0, 1))), 0.0)
        mag = u * w
        with np.errstate(invalid="ignore", divide="ignore"):
            dirv = np.where(d[:, None] > 1e-9, radial / d[:, None], 0.0)
        return dirv * mag[:, None]

    def magnitude(self, points):
        """Displacement magnitude (N, mm)."""
        return np.linalg.norm(self.field(points), axis=1)

    def apply(self, points):
        """Return displaced positions (original + displacement)."""
        P = np.atleast_2d(np.asarray(points, dtype=np.float64))
        return P + self.field(P)

    @property
    def max_shift_mm(self) -> float:
        """Maximum predicted tissue displacement (at the probe wall)."""
        a, R = self.probe_radius_mm, self.influence_radius_mm
        # max of u*w occurs near the axis; at d->0, u->a, w->1
        return float(a)

    def summary(self, n_axis=60, n_ring=16, ring_mm=3.0) -> dict:
        """Summary stats sampled on a thin cylindrical shell around the probe.

        (Radial displacement direction is undefined exactly on the axis, so we
        sample a small ring of radius `ring_mm` to characterise near-probe shift.)
        """
        E, T = self.trajectory.entry, self.trajectory.target
        v = T - E
        L = float(np.linalg.norm(v))
        vhat = v / L if L > 0 else v
        # build an orthonormal frame around the axis
        ref = np.array([1.0, 0, 0]) if abs(vhat[0]) < 0.9 else np.array([0, 1.0, 0])
        e1 = np.cross(vhat, ref); e1 /= np.linalg.norm(e1)
        e2 = np.cross(vhat, e1)
        ts = np.linspace(0, 1, n_axis)
        angles = np.linspace(0, 2 * np.pi, n_ring, endpoint=False)
        pts = []
        for t in ts:
            base = E + t * v
            for ang in angles:
                pts.append(base + ring_mm * (np.cos(ang) * e1 + np.sin(ang) * e2))
        mags = self.magnitude(np.array(pts))
        return {"max_shift_mm": round(self.max_shift_mm, 2),
                f"mean_shift_at_{ring_mm:.0f}mm": round(float(mags.mean()), 2),
                "influence_radius_mm": self.influence_radius_mm}
