"""Multi-trajectory surgical plan: a named collection of trajectories with
optional per-trajectory risk reports (e.g., bilateral DBS leads, multiple SEEG
electrodes)."""
from __future__ import annotations
import json
from dataclasses import dataclass, field
import numpy as np
from .planning import Trajectory


@dataclass
class PlannedTrajectory:
    name: str
    trajectory: Trajectory
    risk: dict = None


@dataclass
class SurgicalPlan:
    patient_id: str = "DEMO"
    note: str = ""
    items: list = field(default_factory=list)

    def add(self, name, trajectory: Trajectory, risk: dict = None) -> "SurgicalPlan":
        self.items.append(PlannedTrajectory(name, trajectory, risk))
        return self

    def __len__(self):
        return len(self.items)

    def __iter__(self):
        return iter(self.items)


def _sanitize(obj):
    """Make a risk dict JSON-safe (numpy -> native, drop arrays that don't fit)."""
    if isinstance(obj, dict):
        return {k: _sanitize(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_sanitize(v) for v in obj]
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, (np.floating, np.integer)):
        return obj.item()
    return obj


def save_plan(plan: SurgicalPlan, path: str) -> str:
    """Serialize a plan to JSON (entry/target in world mm, plus any risk dict).

    Coordinates are stored, not imagery, so a plan is small and portable; reopen
    it against the same patient volume to restore the trajectories."""
    data = {"format": "sulcora-plan/1", "patient_id": plan.patient_id, "note": plan.note,
            "items": [{"name": it.name,
                       "entry": np.asarray(it.trajectory.entry, float).tolist(),
                       "target": np.asarray(it.trajectory.target, float).tolist(),
                       "risk": _sanitize(it.risk) if it.risk else None}
                      for it in plan]}
    with open(path, "w") as f:
        json.dump(data, f, indent=2)
    return path


def load_plan(path: str) -> SurgicalPlan:
    """Load a plan previously written by :func:`save_plan`."""
    with open(path) as f:
        d = json.load(f)
    plan = SurgicalPlan(patient_id=d.get("patient_id", "DEMO"), note=d.get("note", ""))
    for it in d.get("items", []):
        plan.add(it["name"], Trajectory(it["entry"], it["target"]), it.get("risk"))
    return plan
