"""Finite-element tissue-deformation model (linear elasticity).

This upgrades the closed-form ``deformation.ProbeDeformation`` heuristic to an
actual continuum-mechanics solve: small-strain **linear elasticity** discretised
with 8-node trilinear hexahedral finite elements on a regular grid, solving the
Navier-Cauchy equilibrium equations

    div( sigma ) = 0,   sigma = lambda*tr(eps) I + 2 mu eps,   eps = sym(grad u)

for the tissue displacement field ``u`` produced by advancing a probe along the
trajectory. The probe is modelled as a displacement-driven cylindrical cavity
(the wall is pushed radially outward), with the far field (skull) held fixed.

Unlike the analytical model, this respects the elastic coupling between
directions (Poisson effect), the domain geometry, and the fixed outer boundary —
so the predicted shift is a genuine elastostatic solution rather than a per-point
formula. A useful property of this displacement-driven, body-force-free setup:
the solution is **independent of the absolute stiffness** (Young's modulus
cancels), depending only on Poisson's ratio and geometry — so no patient-specific
modulus is needed for the planning preview.

IMPORTANT: this is still a *research* model — homogeneous, isotropic, linear,
small-strain, on the preoperative geometry. It is the engineering stepping stone
toward the validated FEM-trained AI surrogate (verified against measured tissue
displacement) described in docs/VALIDATION_PLAN.md. Not a medical device; not for
clinical use.
"""
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np
from .planning import Trajectory

# Local node ordering for the trilinear hex (natural-coordinate sign of each
# corner) and the matching (di,dj,dk) grid offset.
_SIGNS = np.array([
    [-1, -1, -1], [1, -1, -1], [1, 1, -1], [-1, 1, -1],
    [-1, -1, 1], [1, -1, 1], [1, 1, 1], [-1, 1, 1],
], dtype=float)
_OFFS = ((_SIGNS + 1) / 2).astype(int)          # corner offsets in {0,1}^3
_GP = np.array([-1, 1]) / np.sqrt(3.0)          # 2-point Gauss abscissae


def _elasticity_D(nu: float) -> np.ndarray:
    """Isotropic elastic stiffness (Voigt 6x6) with E factored out (E=1).

    Displacement-driven, load-free solutions are independent of E, so we set
    E=1 and keep only the Poisson dependence.
    """
    E = 1.0
    lam = E * nu / ((1 + nu) * (1 - 2 * nu))
    mu = E / (2 * (1 + nu))
    D = np.zeros((6, 6))
    D[:3, :3] = lam
    D[0, 0] = D[1, 1] = D[2, 2] = lam + 2 * mu
    D[3, 3] = D[4, 4] = D[5, 5] = mu
    return D


def _hex_stiffness(h: float, nu: float) -> np.ndarray:
    """24x24 stiffness of a cubic 8-node hex element of side `h`."""
    D = _elasticity_D(nu)
    Ke = np.zeros((24, 24))
    detJ = (h / 2.0) ** 3
    dxi_dx = 2.0 / h                       # natural -> physical derivative scale
    for gx in _GP:
        for gy in _GP:
            for gz in _GP:
                # shape-function natural derivatives at this Gauss point
                dN = np.empty((8, 3))
                for a, (sx, sy, sz) in enumerate(_SIGNS):
                    dN[a, 0] = 0.125 * sx * (1 + sy * gy) * (1 + sz * gz)
                    dN[a, 1] = 0.125 * sy * (1 + sx * gx) * (1 + sz * gz)
                    dN[a, 2] = 0.125 * sz * (1 + sx * gx) * (1 + sy * gy)
                dN *= dxi_dx
                B = np.zeros((6, 24))
                for a in range(8):
                    bx, by, bz = dN[a]
                    c = 3 * a
                    B[0, c] = bx
                    B[1, c + 1] = by
                    B[2, c + 2] = bz
                    B[3, c] = by; B[3, c + 1] = bx
                    B[4, c + 1] = bz; B[4, c + 2] = by
                    B[5, c] = bz; B[5, c + 2] = bx
                Ke += (B.T @ D @ B) * detJ      # Gauss weights are 1
    return Ke


@dataclass
class FEMProbeDeformation:
    """Linear-elastic FEM model of probe-induced tissue displacement.

    Drop-in compatible with ``ProbeDeformation``: exposes ``field``,
    ``magnitude``, ``apply``, ``max_shift_mm`` and ``summary`` so it can be passed
    straight to ``risk.deformation_aware_risk`` and the robotics simulator.

    trajectory          : the planned probe path (world mm).
    probe_radius_mm     : cavity radius / outward wall displacement.
    influence_radius_mm : padding of the solve domain around the trajectory; the
                          outer boundary is held fixed (skull).
    poisson             : Poisson's ratio (~0.45 for near-incompressible brain).
    element_mm          : FEM element edge length (solve resolution).
    """
    trajectory: Trajectory
    probe_radius_mm: float = 1.5
    influence_radius_mm: float = 18.0
    poisson: float = 0.45
    element_mm: float = 3.0
    _solved: bool = field(default=False, repr=False)

    # ---- geometry helpers ----
    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, :]

    # ---- mesh + solve (lazy, cached) ----
    def _ensure_solved(self):
        if self._solved:
            return
        E, T = self.trajectory.entry, self.trajectory.target
        h = float(self.element_mm)
        pad = self.influence_radius_mm
        lo = np.minimum(E, T) - pad
        hi = np.maximum(E, T) + pad
        # snap the origin so the entry point lands exactly on a node; this aligns
        # the probe axis with grid lines and avoids an off-centre cavity wall
        # (which would make an axis-aligned probe spuriously asymmetric).
        lo = E - np.ceil((E - lo) / h) * h
        dims = np.maximum(np.ceil((hi - lo) / h).astype(int), 1)   # elements/axis
        nx, ny, nz = dims + 1                                       # nodes/axis
        self._origin, self._h, self._ndim = lo, h, (nx, ny, nz)

        # node world coordinates
        gi, gj, gk = np.meshgrid(np.arange(nx), np.arange(ny), np.arange(nz), indexing="ij")
        coords = lo[None, :] + np.stack([gi, gj, gk], axis=-1).reshape(-1, 3) * h
        N = coords.shape[0]

        def nid(i, j, k):
            return (i * ny + j) * nz + k

        # element connectivity (E,8)
        ei, ej, ek = np.meshgrid(np.arange(dims[0]), np.arange(dims[1]),
                                  np.arange(dims[2]), indexing="ij")
        base = np.stack([ei, ej, ek], axis=-1).reshape(-1, 3)
        conn = np.empty((base.shape[0], 8), dtype=np.int64)
        for a, off in enumerate(_OFFS):
            conn[:, a] = nid(base[:, 0] + off[0], base[:, 1] + off[1], base[:, 2] + off[2])

        # element DOF map (E,24) and identical element stiffness
        edof = (conn[:, :, None] * 3 + np.arange(3)[None, None, :]).reshape(conn.shape[0], 24)
        Ke = _hex_stiffness(h, self.poisson)

        from scipy.sparse import coo_matrix
        rows = np.broadcast_to(edof[:, :, None], (edof.shape[0], 24, 24)).reshape(-1)
        cols = np.broadcast_to(edof[:, None, :], (edof.shape[0], 24, 24)).reshape(-1)
        vals = np.broadcast_to(Ke[None, :, :], (edof.shape[0], 24, 24)).reshape(-1)
        K = coo_matrix((vals, (rows, cols)), shape=(3 * N, 3 * N)).tocsr()

        # ---- boundary conditions ----
        closest = self._closest_on_axis(coords)
        radial = coords - closest
        d = np.linalg.norm(radial, axis=1)
        r_in = max(self.probe_radius_mm, h)               # guarantee >=1 wall ring
        inner = d <= r_in
        # outer box faces held fixed (skull / far field)
        onface = ((coords[:, 0] <= lo[0] + 1e-6) | (coords[:, 0] >= hi[0] - 1e-6) |
                  (coords[:, 1] <= lo[1] + 1e-6) | (coords[:, 1] >= hi[1] - 1e-6) |
                  (coords[:, 2] <= lo[2] + 1e-6) | (coords[:, 2] >= hi[2] - 1e-6))

        u = np.zeros((N, 3))
        with np.errstate(invalid="ignore", divide="ignore"):
            rhat = np.where(d[:, None] > 1e-9, radial / d[:, None], 0.0)
        u[inner] = rhat[inner] * self.probe_radius_mm     # wall pushed outward
        # prescribed DOFs = inner wall + fixed faces (faces keep u=0)
        pres_node = inner | onface
        pres_dofs = np.where(np.repeat(pres_node, 3))[0]
        free_dofs = np.setdiff1d(np.arange(3 * N), pres_dofs, assume_unique=False)

        uvec = u.reshape(-1)
        from scipy.sparse.linalg import spsolve
        Kff = K[free_dofs][:, free_dofs]
        Kfp = K[free_dofs][:, pres_dofs]
        rhs = -Kfp.dot(uvec[pres_dofs])
        uf = spsolve(Kff.tocsc(), rhs)
        uvec[free_dofs] = uf

        self._U = uvec.reshape(nx, ny, nz, 3)
        self._solved = True

    # ---- field evaluation (trilinear interpolation of the FEM solution) ----
    def field(self, points):
        """Displacement vectors (N,3 mm) for input world points (N,3)."""
        self._ensure_solved()
        P = np.atleast_2d(np.asarray(points, dtype=np.float64))
        nx, ny, nz = self._ndim
        g = (P - self._origin[None, :]) / self._h            # continuous grid coords
        out = np.zeros((P.shape[0], 3))
        i0 = np.floor(g).astype(int)
        inside = np.all((i0 >= 0) & (i0 <= np.array([nx - 2, ny - 2, nz - 2])), axis=1)
        if not np.any(inside):
            return out
        idx = np.where(inside)[0]
        i = i0[idx]
        fr = g[idx] - i
        U = self._U
        for n in range(idx.shape[0]):
            ii, jj, kk = i[n]
            fx, fy, fz = fr[n]
            c = U[ii:ii + 2, jj:jj + 2, kk:kk + 2]           # (2,2,2,3)
            c00 = c[0, 0, 0] * (1 - fx) + c[1, 0, 0] * fx
            c01 = c[0, 0, 1] * (1 - fx) + c[1, 0, 1] * fx
            c10 = c[0, 1, 0] * (1 - fx) + c[1, 1, 0] * fx
            c11 = c[0, 1, 1] * (1 - fx) + c[1, 1, 1] * fx
            c0 = c00 * (1 - fy) + c10 * fy
            c1 = c01 * (1 - fy) + c11 * fy
            out[idx[n]] = c0 * (1 - fz) + c1 * fz
        return out

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

    def apply(self, points):
        P = np.atleast_2d(np.asarray(points, dtype=np.float64))
        return P + self.field(P)

    @property
    def max_shift_mm(self) -> float:
        """Maximum predicted displacement (the prescribed cavity-wall motion)."""
        return float(self.probe_radius_mm)

    def summary(self, n_axis=60, n_ring=16, ring_mm=3.0) -> dict:
        """Displacement stats sampled on a thin cylindrical shell around the probe."""
        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
        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)
        pts = []
        for t in np.linspace(0, 1, n_axis):
            base = E + t * v
            for ang in np.linspace(0, 2 * np.pi, n_ring, endpoint=False):
                pts.append(base + ring_mm * (np.cos(ang) * e1 + np.sin(ang) * e2))
        mags = self.magnitude(np.array(pts))
        return {"model": "FEM linear elasticity",
                "max_shift_mm": round(self.max_shift_mm, 2),
                f"mean_shift_at_{ring_mm:.0f}mm": round(float(mags.mean()), 2),
                "poisson": self.poisson,
                "element_mm": self.element_mm,
                "influence_radius_mm": self.influence_radius_mm}
