import os
import torch
import SimpleITK as sitk

# --- Download / locate the trained model weights (~400MB) ---
# License reminder: The official nnInteractive checkpoint is licensed under
# Creative Commons Attribution Non Commercial Share Alike 4.0 (CC BY-NC-SA 4.0).
# See the License section of this readme!
#
# nnInteractive ships a small model-management API. It fetches the list of available
# models from Hugging Face (MIC-DKFZ/nnInteractive), downloads only the selected one,
# reuses it on later runs, and works offline once a model has been downloaded. Models
# are stored under $NNINTERACTIVE_MODEL_DIR (default: ~/.nninteractive).
from nnInteractive.model_management import ensure_model_available, get_default_model_id

# Use the recommended model — i.e. the manifest's default. Don't hard-code a version here, so you
# automatically track whatever is currently recommended:
model_id = get_default_model_id()              # resolves to the recommended id, e.g. "nnInteractive_v1.0"
# To pin a specific model instead, set its id by name (see the ids from
# `nninteractive-available-models` or nnInteractive.model_management.list_models()):
#   model_id = "nnInteractive_v1.0"

# Resolve the id to a local folder: downloads on first use, reuses it afterwards (offline-friendly).
model_path = ensure_model_available(model_id)
# ...later passed to the session via session.initialize_from_trained_model_folder(str(model_path)).

# --- Initialize Inference Session ---
from nnInteractive.inference.inference_session import nnInteractiveInferenceSession

session = nnInteractiveInferenceSession(
    device=torch.device("cuda:0"),  # Set inference device
    use_torch_compile=False,  # Compiles the net for faster predictions. The one-time (slow) compile
                              # is paid during initialize_from_trained_model_folder() below (it warms
                              # up automatically), NOT on your first prompt. Worth it for long-lived
                              # processes (the server enables it by default) or longer sessions.
    verbose=False,
    torch_n_threads=os.cpu_count(),  # Use available CPU cores, cap this if your system has a gigantic CPU (64+ cores)
    do_autozoom=True,  # Enables AutoZoom for better patching
)

# Load the trained model
session.initialize_from_trained_model_folder(str(model_path))

# --- Load Input Image (Example with SimpleITK) ---
# DO NOT preprocess the image in any way. Give it to nnInteractive as it is! DO NOT apply level window, DO NOT normalize 
# intensities and never ever convert an image with higher precision (float32, uint16, etc) to uint8!
# The ONLY instance where some preprocesing makes sense is if your original image is too large to be reasonably used. 
# This may be the case, for example, for some microCT images. In this case you can consider downsampling.
input_image = sitk.ReadImage("FILENAME")
img = sitk.GetArrayFromImage(input_image)[None]  # Ensure shape (1, x, y, z)

# Validate input dimensions
if img.ndim != 4:
    raise ValueError("Input image must be 4D with shape (1, x, y, z)")

session.set_image(img)

# --- Define Output Buffer ---
target_tensor = torch.zeros(img.shape[1:], dtype=torch.uint8)  # Must be 3D (x, y, z)
session.set_target_buffer(target_tensor)

# --- Interacting with the Model ---
# Interactions can be freely chained and mixed in any order. Each interaction refines the
# segmentation, and the model writes the updated mask straight into your target buffer
# (in place) after every interaction — you never have to fetch anything back.
#
# Normally you just read the target buffer directly (see "Retrieve Results" below). The only
# reason to look at the return value is if you must propagate the update into your OWN structure
# (a GUI's label layer, a separate array, ...) and want to avoid copying the whole volume every
# time: each add_*_interaction(..., run_prediction=True) RETURNS the bounding box of the region it
# changed, as [[x1, x2], [y1, y2], [z1, z2]] in target-buffer coordinates (None if nothing changed
# / run_prediction=False), so you can copy just that sub-volume.

# Example: Add a **positive** point interaction
# POINT_COORDINATES should be a tuple (x, y, z) specifying the point location.
changed_bbox = session.add_point_interaction(POINT_COORDINATES, include_interaction=True)
# Only needed if you mirror the result elsewhere (otherwise ignore it and read target_tensor):
# if changed_bbox is not None:
#     (x1, x2), (y1, y2), (z1, z2) = changed_bbox
#     my_label_volume[x1:x2, y1:y2, z1:z2] = target_tensor[x1:x2, y1:y2, z1:z2]

# Example: Add a **negative** point interaction
# To make any interaction negative set include_interaction=False
session.add_point_interaction(POINT_COORDINATES, include_interaction=False)

# Example: Add a bounding box interaction
# BBOX_COORDINATES must be specified as [[x1, x2], [y1, y2], [z1, z2]] (half-open intervals).
# Note: nnInteractive pre-trained models currently only support **2D bounding boxes**.
# This means that **one dimension must be [d, d+1]** to indicate a single slice.

# Example of a 2D bounding box in the axial plane (XY slice at depth Z)
# BBOX_COORDINATES = [[30, 80], [40, 100], [10, 11]]  # X: 30-80, Y: 40-100, Z: slice 10

session.add_bbox_interaction(BBOX_COORDINATES, include_interaction=True)

# Example: Add a scribble interaction
# - Background must be 0, and scribble must be 1.
# - Use session.preferred_scribble_thickness for optimal results.
#
# ✅ RECOMMENDED (v2): pass a small 2D crop plus its location.
# Scribbles live on a single axis-aligned slice, so one of the three bbox
# dimensions is always size 1 and the in-plane extent typically covers only
# a small region. The cropped array is ORDERS OF MAGNITUDE
# smaller than a full-volume mask for typical annotations, which makes this
# path dramatically faster. Please prefer this
# form in new integrations.
#
# SCRIBBLE_CROP.shape must equal the bbox size. INTERACTION_BBOX uses
# half-open intervals [[x1,x2],[y1,y2],[z1,z2]] in original-image coordinates.
# Example: a scribble drawn on axial slice z=64, covering x∈[100,140), y∈[80,150):
#   SCRIBBLE_CROP    = <ndarray of shape (40, 70, 1), values 0 or 1>
#   INTERACTION_BBOX = [[100, 140], [80, 150], [64, 65]]
session.add_scribble_interaction(
    SCRIBBLE_CROP,
    include_interaction=True,
    interaction_bbox=INTERACTION_BBOX,
)

# Legacy form (still supported, but discouraged): a 3D array matching the
# full original image shape with the scribble baked into one slice.
# session.add_scribble_interaction(SCRIBBLE_IMAGE, include_interaction=True)

# Example: Add a lasso interaction
# - Like scribble but the single slice contains a **closed contour** for the selection.
# - Same recommendation applies: pass a 2D crop + interaction_bbox for a large speedup.
session.add_lasso_interaction(
    LASSO_CROP,
    include_interaction=True,
    interaction_bbox=INTERACTION_BBOX,
)
# Legacy full-volume form (discouraged):
# session.add_lasso_interaction(LASSO_IMAGE, include_interaction=True)

# You can combine any number of interactions as needed. 
# The model refines the segmentation result incrementally with each new interaction.

# --- Retrieve Results ---
# The result already lives in your target buffer: target_tensor IS session.target_buffer (the
# same object, written in place). So just read it — no copy needed to inspect it or save it:
result_np = target_tensor.cpu().numpy()
# sitk.WriteImage(sitk.GetImageFromArray(result_np), "segmentation.nii.gz")

# You only need a COPY if you want to keep this result in memory while you reuse the buffer for
# the next object, because reset_interactions() / reusing the buffer overwrites it in place:
# saved = target_tensor.clone()        # torch  (numpy buffer: target_tensor.copy())

# --- Start a New Object Segmentation ---
session.reset_interactions()  # Clears the target buffer and resets interactions
# (Alternatively, give each object its own fresh buffer instead of resetting:)
# session.set_target_buffer(torch.zeros(img.shape[1:], dtype=torch.uint8))

# Now you can start segmenting the next object in the image.

# --- Set a New Image ---
# Setting a new image also requires setting a new matching target buffer
session.set_image(NEW_IMAGE)
session.set_target_buffer(torch.zeros(NEW_IMAGE.shape[1:], dtype=torch.uint8))

# Enjoy!

from nnInteractive.inference.remote import nnInteractiveRemoteInferenceSession

# Requires the server above to be running and reachable at this URL.
session = nnInteractiveRemoteInferenceSession(
    server_url="http://gpu-box.lab:1527",
    api_key="…",
)
# From here on, the API is identical to nnInteractiveInferenceSession
# (set_image / set_target_buffer / add_*_interaction / ...).