"""Demonstration: anisotropic deformation/error-aware margin vs. a fixed margin.

Shows the non-obvious result: two corridors whose structures sit at the SAME static
clearance, but one structure is predicted to deform TOWARD the corridor and the
other AWAY. A conventional fixed 2 mm margin cannot tell them apart (both "clear"),
so it may pick the dangerous corridor. The anisotropic required margin correctly
flags the toward-deforming case as unsafe and selects the away-deforming corridor.
This is the evidence that the method solves the "arbitrary fixed margin" problem
the literature itself criticizes.

RESEARCH SOFTWARE — NOT FOR CLINICAL USE.
"""
import os, sys
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, HERE)
from sulcora.volume import Volume
from sulcora.planning import Trajectory
from sulcora.optimizer import RiskStructure
from sulcora import margin as M

FIXED = 2.0  # conventional fixed standoff (mm)


def _mask_at(center, shape=(40, 40, 50)):
    m = np.zeros(shape, np.uint8); x, y, z = center; m[x-1:x+2, y-1:y+2, z-1:z+2] = 1
    return Volume(m, np.eye(4))


def main():
    docs = os.path.join(HERE, "docs", "images"); os.makedirs(docs, exist_ok=True)
    # two corridors along z at x=20; each has a structure ~3 mm laterally (at x=23)
    trajA = Trajectory(np.array([20, 18, 48.0]), np.array([20, 18, 6.0]))
    trajB = Trajectory(np.array([20, 22, 48.0]), np.array([20, 22, 6.0]))
    structA = RiskStructure("A_vessel", _mask_at((23, 18, 25)), 2.0, 1.0)   # 3 mm from corridor A
    structB = RiskStructure("B_vessel", _mask_at((23, 22, 25)), 2.0, 1.0)   # 3 mm from corridor B
    toward = lambda p: p + np.array([-2.5, 0, 0])     # A's structure pushed toward corridor
    away = lambda p: p + np.array([2.5, 0, 0])        # B's structure pushed away

    mA = M.corridor_margins(trajA, [structA], deform=toward, robot_sigma_mm=1.0, robot_sigma_deg=1.0)
    mB = M.corridor_margins(trajB, [structB], deform=away, robot_sigma_mm=1.0, robot_sigma_deg=1.0)
    rA, rB = mA["per_structure"]["A_vessel"], mB["per_structure"]["B_vessel"]

    print(f"{'':10}{'clearance':>10}{'fixed 2mm':>12}{'toward':>9}{'required':>10}{'feasible':>10}")
    for name, r in [("Corridor A", rA), ("Corridor B", rB)]:
        fixed_ok = r["actual_clearance_mm"] >= FIXED
        print(f"{name:10}{r['actual_clearance_mm']:>10}{('SAFE' if fixed_ok else 'UNSAFE'):>12}"
              f"{r['toward_mm']:>9}{r['required_mm']:>10}{('YES' if r['feasible'] else 'NO'):>10}")

    fixed_pick = "A or B (indistinguishable)" if (rA['actual_clearance_mm'] >= FIXED and rB['actual_clearance_mm'] >= FIXED) else "?"
    aniso_pick = "B" if (rB["feasible"] and not rA["feasible"]) else ("A" if rA["feasible"] and not rB["feasible"] else "?")
    print(f"\nFixed 2 mm margin picks: {fixed_pick}")
    print(f"Anisotropic required margin picks: Corridor {aniso_pick}  (A flagged UNSAFE: structure deforms toward it)")

    # figure
    fig, ax = plt.subplots(1, 2, figsize=(11, 5), facecolor="#0b0f18")
    for a, (name, r, tw) in zip(ax, [("Corridor A — structure deforms TOWARD", rA, -2.5),
                                     ("Corridor B — structure deforms AWAY", rB, 2.5)]):
        a.set_facecolor("#0b0f18"); a.set_xlim(-8, 8); a.set_ylim(-8, 8); a.set_aspect("equal")
        a.axvline(0, color="#5db1ff", lw=2)                          # corridor (cross-section)
        sx = 3 + tw                                                  # structure center after deformation
        a.add_patch(plt.Circle((3, 0), 1.0, ec="#39ff9e", fc="none", lw=1.4))         # static structure
        a.add_patch(plt.Circle((sx, 0), 1.0, ec="#ff5d5d", fc="none", lw=1.8))        # deformed
        a.add_patch(plt.Circle((0, 0), FIXED, ec="#f5b041", fc="none", ls=":", lw=1.4))   # fixed margin
        a.add_patch(plt.Circle((0, 0), r["required_mm"], ec="#ff5d5d" if not r["feasible"] else "#2ecc71",
                               fc="none", lw=1.8))                   # required margin
        a.set_title(name, color="#9fd0ff", fontsize=10)
        a.text(-7.5, -7, f"clear 3.0 | fixed2mm {'SAFE' if r['actual_clearance_mm']>=FIXED else 'UNSAFE'} | "
               f"required {r['required_mm']}mm | {'FEASIBLE' if r['feasible'] else 'UNSAFE'}",
               color="#e8f0ff", fontsize=8)
        a.axis("off")
    fig.suptitle("Fixed 2 mm margin can't distinguish A from B; anisotropic required margin flags A unsafe",
                 color="#dce6f5", fontsize=12)
    fig.text(0.5, 0.02, "Dotted amber = fixed 2 mm.  Solid = anisotropic required margin (red=unsafe, green=feasible).  "
             "RESEARCH SOFTWARE — NOT FOR CLINICAL USE.", color="#c0392b", ha="center", fontsize=8)
    out = os.path.join(docs, "required_margin_demo.png")
    fig.tight_layout(rect=[0, 0.05, 1, 0.94]); fig.savefig(out, dpi=130, facecolor=fig.get_facecolor())
    print("\nwrote", out)
    assert not rA["feasible"] and rB["feasible"], "demo should flag A unsafe, B feasible"
    print("MARGIN DEMONSTRATION OK — fixed margin and anisotropic margin disagree, as intended")


if __name__ == "__main__":
    main()
