"""Deformation- and error-aware required safety margin (anisotropic, per-structure).

Conventional planners use a fixed, arbitrary standoff (e.g., 2-3 mm) around every
critical structure. This module instead computes, for a given corridor and a
given robotic system's targeting-error specification, the *required* standoff for
each critical structure — and it is **anisotropic**: a structure predicted to
deform TOWARD the corridor (pushed in by the advancing probe) demands more
margin, while one deforming AWAY demands less. The required standoff is obtained
by inverting a contact-probability model that combines (a) the robot's error
distribution and (b) the toward-corridor component of the predicted probe-induced
displacement field, to emit the minimum clearance guaranteeing at most a target
contact probability. Corridors are then ranked by this calibrated margin.

This turns the deformation model + robot-error bench from a validation output into
a planning input. Pure NumPy/SciPy; unit-tested. Research software; not a medical device.
"""
from __future__ import annotations
import numpy as np
from .vessels import _mask_world_points


def _lateral_sigma(depth_mm, sigma_mm, sigma_deg):
    """Lateral positional 1-sigma at a given depth: entry placement error combined
    with angular error accumulated over depth."""
    return float(np.hypot(sigma_mm, np.deg2rad(sigma_deg) * depth_mm))


def _normal_ppf(p):
    """Inverse standard-normal CDF (Acklam approximation) — avoids a scipy.stats dep."""
    a = [-3.969683028665376e+01, 2.209460984245205e+02, -2.759285104469687e+02,
         1.383577518672690e+02, -3.066479806614716e+01, 2.506628277459239e+00]
    b = [-5.447609879822406e+01, 1.615858368580409e+02, -1.556989798598866e+02,
         6.680131188771972e+01, -1.328068155288572e+01]
    c = [-7.784894002430293e-03, -3.223964580411365e-01, -2.400758277161838e+00,
         -2.549732539343734e+00, 4.374664141464968e+00, 2.938163982698783e+00]
    d = [7.784695709041462e-03, 3.224671290700398e-01, 2.445134137142996e+00,
         3.754408661907416e+00]
    plow, phigh = 0.02425, 1 - 0.02425
    if p < plow:
        q = np.sqrt(-2 * np.log(p))
        return (((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) / ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1)
    if p > phigh:
        q = np.sqrt(-2 * np.log(1 - p))
        return -(((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) / ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1)
    q = p - 0.5; r = q * q
    return (((((a[0]*r+a[1])*r+a[2])*r+a[3])*r+a[4])*r+a[5])*q / (((((b[0]*r+b[1])*r+b[2])*r+b[3])*r+b[4])*r+1)


def _normal_sf(x):
    """Standard-normal survival function 1-Phi(x) via erfc."""
    from math import erfc, sqrt
    return 0.5 * erfc(x / sqrt(2.0))


def required_standoff(traj, structure_mask, deform=None, robot_sigma_mm=1.0,
                      robot_sigma_deg=1.0, contact_prob=0.01, n=240) -> dict:
    """Required standoff (mm) for one structure on this corridor and this robot.

    `deform` maps world points (M,3) -> displaced points (M,3) for the predicted
    probe-induced field of this corridor (e.g., from `deformation`/`fem`); None =
    no deformation. Returns the required margin, the toward-corridor displacement
    at the binding point, the robot lateral sigma there, the static clearance, the
    contact probability at that clearance, and feasibility.
    """
    orig = _mask_world_points(structure_mask)
    if len(orig) == 0:
        return None
    disp = np.asarray(deform(orig), float) if deform is not None else orig
    dvec = disp - orig                                   # displacement vectors
    pts = traj.sample(n)
    from scipy.spatial import cKDTree
    dist, idx = cKDTree(pts).query(orig)                 # nearest corridor sample per point
    u = orig - pts[idx]                                  # radial: corridor -> structure
    un = np.linalg.norm(u, axis=1, keepdims=True); un[un == 0] = 1.0
    u = u / un
    toward = -np.sum(dvec * u, axis=1)                   # + = structure moves toward corridor
    depth = (idx.astype(float) / max(n - 1, 1)) * traj.length
    sig = np.hypot(robot_sigma_mm, np.deg2rad(robot_sigma_deg) * depth)
    z = _normal_ppf(1.0 - contact_prob)
    required = toward + z * sig                          # per-point required standoff (anisotropic)
    # binding point: smallest safety headroom (static distance minus required standoff)
    b = int(np.argmin(dist - required))
    req_b = float(required[b]); actual_b = float(dist[b])
    cp = _normal_sf((actual_b - toward[b]) / max(sig[b], 1e-6))
    return {"required_mm": round(req_b, 2), "actual_clearance_mm": round(actual_b, 2),
            "toward_mm": round(float(toward[b]), 2), "sigma_lat_mm": round(float(sig[b]), 2),
            "contact_prob": round(float(cp), 4), "feasible": actual_b >= req_b}


def corridor_margins(traj, structures, deform=None, robot_sigma_mm=1.0,
                     robot_sigma_deg=1.0, contact_prob=0.01, n=240) -> dict:
    """Per-structure required standoffs for a corridor plus the binding structure.

    `structures` is a list of objects with `.name` and `.mask` (e.g., RiskStructure).
    Returns per-structure results, the binding (most-demanding) structure, the
    overall required margin, and whether the corridor is feasible for this robot.
    """
    per = {}
    for s in structures:
        r = required_standoff(traj, s.mask, deform=deform, robot_sigma_mm=robot_sigma_mm,
                              robot_sigma_deg=robot_sigma_deg, contact_prob=contact_prob, n=n)
        if r is not None:
            per[s.name] = r
    if not per:
        return {"per_structure": {}, "binding": None, "required_mm": 0.0, "feasible": True}
    binding = max(per, key=lambda k: per[k]["required_mm"] - per[k]["actual_clearance_mm"])
    return {"per_structure": per, "binding": binding,
            "required_mm": per[binding]["required_mm"],
            "feasible": all(v["feasible"] for v in per.values())}


def standoff_envelope(traj, structure_mask, deform=None, robot_sigma_mm=1.0,
                      robot_sigma_deg=1.0, contact_prob=0.01, n=240):
    """Anisotropic keep-out envelope: the required standoff at EACH surface point of
    the structure, shaped by the predicted deformation direction. Returns (points,
    required) where required[i] is the standoff for points[i] — larger where the
    surface is predicted to move toward the corridor. This is the per-surface-point
    (non-spherical) form of the required margin."""
    orig = _mask_world_points(structure_mask)
    if len(orig) == 0:
        return np.empty((0, 3)), np.empty((0,))
    disp = np.asarray(deform(orig), float) if deform is not None else orig
    dvec = disp - orig
    pts = traj.sample(n)
    from scipy.spatial import cKDTree
    dist, idx = cKDTree(pts).query(orig)
    u = orig - pts[idx]; un = np.linalg.norm(u, axis=1, keepdims=True); un[un == 0] = 1.0; u = u / un
    toward = -np.sum(dvec * u, axis=1)
    depth = (idx.astype(float) / max(n - 1, 1)) * traj.length
    sig = np.hypot(robot_sigma_mm, np.deg2rad(robot_sigma_deg) * depth)
    required = toward + _normal_ppf(1.0 - contact_prob) * sig
    return orig, required


def required_device_accuracy(traj, structures, deform=None, contact_prob=0.01,
                             ref_deg_per_mm=1.0, n=240) -> dict:
    """INVERSE: solve for the maximum device targeting error (lateral 1-sigma, mm)
    permissible on this corridor so that P(contact any critical structure) stays at
    or below `contact_prob`. Emits a required device-accuracy specification for
    case-based qualification, driven by the predicted-deformation contact constraint
    (not by robot reachability). Returns the allowable lateral sigma, the binding
    structure, and an equivalent (mm, deg) spec using `ref_deg_per_mm`."""
    z = _normal_ppf(1.0 - contact_prob)
    from scipy.spatial import cKDTree
    pts = traj.sample(n)
    worst = None
    for s in structures:
        orig = _mask_world_points(s.mask)
        if len(orig) == 0:
            continue
        disp = np.asarray(deform(orig), float) if deform is not None else orig
        dvec = disp - orig
        dist, idx = cKDTree(pts).query(orig)
        u = orig - pts[idx]; un = np.linalg.norm(u, axis=1, keepdims=True); un[un == 0] = 1.0; u = u / un
        toward = -np.sum(dvec * u, axis=1)
        # for each surface point: max sigma s.t. toward + z*sigma <= clearance
        allow = (dist - toward) / max(z, 1e-6)
        k = int(np.argmin(allow))                 # most-constraining point
        if worst is None or allow[k] < worst[1]:
            worst = (s.name, float(allow[k]))
    if worst is None:
        return {"max_lateral_sigma_mm": float("inf"), "binding": None, "feasible": True}
    name, sig = worst
    sig = max(sig, 0.0)
    return {"max_lateral_sigma_mm": round(sig, 2), "binding": name,
            "required_spec": f"<= {round(sig,2)} mm (or <= {round(sig*ref_deg_per_mm,2)} deg) lateral 1-sigma",
            "feasible": sig > 0.0}


def safety_certificate(traj, structures, deform=None, robot_sigma_mm=1.0,
                       robot_sigma_deg=1.0, threshold=0.01, n=240) -> dict:
    """Per-plan probabilistic safety CERTIFICATE: a computed upper bound on the
    probability of contacting ANY critical structure for THIS corridor and THIS
    device, combining the targeting-error distribution with the predicted
    deformation. Emits a union-bound P(any contact), the per-structure
    probabilities, the binding structure, and a pass/fail against `threshold`."""
    from scipy.spatial import cKDTree
    pts = traj.sample(n)
    per = {}
    for s in structures:
        orig = _mask_world_points(s.mask)
        if len(orig) == 0:
            continue
        disp = np.asarray(deform(orig), float) if deform is not None else orig
        dvec = disp - orig
        dist, idx = cKDTree(pts).query(orig)
        u = orig - pts[idx]; un = np.linalg.norm(u, axis=1, keepdims=True); un[un == 0] = 1.0; u = u / un
        toward = -np.sum(dvec * u, axis=1)
        depth = (idx.astype(float) / max(n - 1, 1)) * traj.length
        sig = np.hypot(robot_sigma_mm, np.deg2rad(robot_sigma_deg) * depth)
        p = np.array([_normal_sf((dist[i] - toward[i]) / max(sig[i], 1e-6)) for i in range(len(dist))])
        per[s.name] = round(float(p.max()), 5)                # binding surface point of this structure
    if not per:
        return {"p_contact_any_bound": 0.0, "per_structure": {}, "binding": None,
                "threshold": threshold, "pass": True}
    union = min(1.0, float(sum(per.values())))                # conservative union bound
    binding = max(per, key=per.get)
    return {"p_contact_any_bound": round(union, 5), "per_structure": per,
            "binding": binding, "threshold": threshold, "pass": union <= threshold}


def rank_by_required_margin(trajectories, structures, deform_factory=None, **kw):
    """Rank corridors by calibrated safety headroom (feasible first, then largest
    margin between static clearance and required standoff at the binding structure).
    `deform_factory(traj)` returns the per-corridor displacement callable (or None)."""
    scored = []
    for tr in trajectories:
        deform = deform_factory(tr) if deform_factory is not None else None
        m = corridor_margins(tr, structures, deform=deform, **kw)
        headroom = min((v["actual_clearance_mm"] - v["required_mm"]) for v in m["per_structure"].values()) \
            if m["per_structure"] else float("inf")
        scored.append((tr, m, headroom))
    scored.sort(key=lambda t: (t[1]["feasible"], t[2]), reverse=True)
    return scored
