"""Robotics-validation environment.

Treats the patient-specific virtual brain as a test bench for a neurosurgical
robot: given a *planned* trajectory, simulate repeated robotic executions with
realistic entry-placement and angular error, and score each run on (a) targeting
accuracy — the distance between the achieved tip and the planned target (TRE),
and (b) safety — the minimum clearance of the achieved path to critical
structures, optionally under predicted tissue deformation.

This is the "for the Development of Neurosurgical Robotics" piece of ARPA-H
Topic 7: a virtual environment to develop and validate robotic targeting before
any contact with a patient. The execution-error model here is a simplified,
configurable placeholder; the funded work replaces it with measured robot and
tissue behaviour and validated tissue deformation.
"""
from __future__ import annotations
from dataclasses import dataclass
import numpy as np
from .planning import Trajectory


def _perp_basis(u):
    ref = np.array([1.0, 0, 0]) if abs(u[0]) < 0.9 else np.array([0, 1.0, 0])
    e1 = np.cross(u, ref); e1 /= np.linalg.norm(e1)
    e2 = np.cross(u, e1)
    return e1, e2


@dataclass
class RoboticExecution:
    planned: Trajectory
    targeting_sigma_mm: float = 1.0      # lateral entry placement error (1 SD)
    angular_sigma_deg: float = 1.0       # angular error (1 SD)

    def simulate(self, n=1000, clearance_fn=None, tre_tol_mm=2.0,
                 clearance_min_mm=3.0, seed=0) -> dict:
        """Run n simulated executions. Returns accuracy + safety statistics.

        clearance_fn: optional callable(Trajectory)->mm giving the achieved
        path's minimum clearance to critical structures (e.g., deformation-aware).
        """
        rng = np.random.default_rng(seed)
        E0, T0 = self.planned.entry, self.planned.target
        v = T0 - E0
        depth = float(np.linalg.norm(v))
        u = v / depth
        e1, e2 = _perp_basis(u)

        tre = np.empty(n)
        offsets = np.empty((n, 2))
        clears = np.empty(n) if clearance_fn else None
        for i in range(n):
            dE = rng.normal(0, self.targeting_sigma_mm) * e1 + rng.normal(0, self.targeting_sigma_mm) * e2
            da1 = np.deg2rad(rng.normal(0, self.angular_sigma_deg))
            da2 = np.deg2rad(rng.normal(0, self.angular_sigma_deg))
            u2 = u + da1 * e1 + da2 * e2
            u2 /= np.linalg.norm(u2)
            E = E0 + dE
            achieved_T = E + u2 * depth
            err = achieved_T - T0
            tre[i] = np.linalg.norm(err)
            offsets[i] = [float(err @ e1), float(err @ e2)]
            if clearance_fn:
                clears[i] = clearance_fn(Trajectory(E, achieved_T))

        out = {
            "n": n,
            "mean_tre_mm": round(float(tre.mean()), 2),
            "p95_tre_mm": round(float(np.percentile(tre, 95)), 2),
            "max_tre_mm": round(float(tre.max()), 2),
            "tre_within_tol_pct": round(float(np.mean(tre <= tre_tol_mm) * 100), 1),
            "tre_tol_mm": tre_tol_mm,
            "_offsets": offsets,
            "_tre": tre,
        }
        if clearance_fn:
            safe = (tre <= tre_tol_mm) & (clears >= clearance_min_mm)
            out.update({
                "mean_clearance_mm": round(float(clears.mean()), 2),
                "min_clearance_mm": round(float(clears.min()), 2),
                "clearance_min_mm": clearance_min_mm,
                "pass_pct": round(float(np.mean(safe) * 100), 1),
                "_clears": clears,
            })
        return out
