# This file was extracted from the HV SDK Docusaurus examples.
# It is intended as a downloadable, runnable companion to the documentation.
# Set HSI_EXAMPLE_BASE_DIR and related env vars to use your own data.
# Source page: /hsi/hv_sdk/examples/basics#plot-rois-on-an-image

# region: setup
import os
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import qtec_hv_sdk as hs
import qtec_hv_sdk.annotations
from qtec_hv_sdk.preprocessing import make_reference
from qtec_hv_sdk.preprocessing import reflectance_calibration


BASE_DIR = Path(os.environ.get("HSI_EXAMPLE_BASE_DIR", "/path/to/HSI_data/nuts"))
if not BASE_DIR.exists():
    raise SystemExit(
        "Run: 'export HSI_EXAMPLE_BASE_DIR=/path/to/HSI_data/' to setup the "
        "folder containing the example datacubes."
    )
CUBE = os.environ.get("HSI_EXAMPLE_CUBE", "mix1.pam")
DARK_REF = os.environ.get("HSI_EXAMPLE_DARK_REF", "dark_ref.pam")
WHITE_REF = os.environ.get("HSI_EXAMPLE_WHITE_REF", "white_ref.pam")
ANNOTATIONS = os.environ.get("HSI_EXAMPLE_ANNOTATIONS", "mix1_roi.json")

WAVELENGTH_MIN_NM = 430.0
WAVELENGTH_MAX_NM = 1700.0


def make_references():
    dark = hs.open(str(BASE_DIR / DARK_REF))
    white = hs.open(str(BASE_DIR / WHITE_REF))
    return make_reference(dark), make_reference(white)


def add_camera_wavelengths(img):
    # Example PAM cubes may not include wavelengths; attach them so downstream
    # wavelength-based code can use img.meta.spectral.wavelengths.
    if img.meta.spectral.wavelengths is not None:
        return img

    meta = img.meta
    meta.spectral = hs.SpectralMeta(
        wavelengths=hs.WavelengthMeta(
            np.linspace(WAVELENGTH_MIN_NM, WAVELENGTH_MAX_NM, img.meta.shape.bands),
            hs.WavelengthUnit.Nanometer,
        )
    )
    return img.with_meta(meta)


def open_reflectance_cube(cube_name=CUBE):
    dark_ref, white_ref = make_references()
    img = add_camera_wavelengths(hs.open(str(BASE_DIR / cube_name)))
    return reflectance_calibration(img, white_ref, dark_ref, clip=True)


def wavelengths_for_image(img):
    wavelengths = img.meta.spectral.wavelengths
    if wavelengths is None:
        raise ValueError("This image does not contain wavelength metadata.")
    if wavelengths.unit != hs.WavelengthUnit.Nanometer:
        raise ValueError(f"Expected wavelengths in nm, got {wavelengths.unit}.")
    return np.array(wavelengths.band_data)


def band_index_for_wavelength(wavelength_nm, img):
    wavelengths = wavelengths_for_image(img)
    return int(np.argmin(np.abs(wavelengths - wavelength_nm)))


def contrast_stretch(image, percentiles=(1, 99)):
    low, high = np.percentile(image, percentiles)
    return np.clip((image - low) / (high - low + 1e-8), 0, 1)


def load_annotations():
    annotations_path = Path(ANNOTATIONS)
    if not annotations_path.is_absolute():
        annotations_path = BASE_DIR / annotations_path
    if not annotations_path.exists():
        raise SystemExit("Set HSI_EXAMPLE_ANNOTATIONS to an annotations JSON file.")
    return hs.annotations.open(str(annotations_path))


def annotation_value(value):
    if value is None or isinstance(value, (str, int, float)):
        return value
    try:
        return value[0]
    except TypeError:
        return value
# end region

# region: example
import matplotlib.patches as mpatches


def patch_top_center(patch):
    """Return (x_center, y_top) in data coordinates for any patch type.

    In image coordinates y increases downward, so the visual top of a shape
    is its minimum y value.
    """
    if isinstance(patch, mpatches.Ellipse):
        x = patch.center[0]
        y = patch.center[1] - patch.height / 2
    elif isinstance(patch, mpatches.Rectangle):
        x = patch.get_x() + patch.get_width() / 2
        y = patch.get_y()
    else:  # Polygon or any other patch — use path vertices (already in data coords)
        verts = patch.get_path().vertices
        x = (verts[:, 0].min() + verts[:, 0].max()) / 2
        y = verts[:, 1].min()
    return x, y


reflectance = open_reflectance_cube()
ann_file = load_annotations()

target_wavelengths = (650, 550, 460)
band_indices = [
    band_index_for_wavelength(wv, reflectance)
    for wv in target_wavelengths
]

sorted_band_indices = sorted(band_indices)
selected_cube = reflectance.select_bands(sorted_band_indices).to_numpy_with_interleave(hs.bip)

rgb = np.empty((*selected_cube.shape[:2], 3), dtype=np.float32)
for channel, band_index in enumerate(band_indices):
    # select_bands() expects increasing band indices, so map each requested
    # RGB wavelength back to its position in the sorted cube.
    selected_channel = sorted_band_indices.index(band_index)
    rgb[:, :, channel] = contrast_stretch(selected_cube[:, :, selected_channel])

fig, ax, patches = hs.annotations.plot_image_with_annotations(
    rgb,
    ann_file,
    label=False,
    show_axes=False,
)

for annot, patch in zip(ann_file.annotations, patches):
    # plot_image_with_annotations() draws the ROI geometry. The returned
    # matplotlib patches let us add our own labels without rebuilding the ROIs.
    label = annotation_value(annot.properties.get("type")) or ""
    x, y = patch_top_center(patch)
    # va="bottom" places the text's baseline at y, so the text body
    # extends toward smaller y values — visually above the ROI.
    ax.text(
        x, y,
        label,
        ha="center",
        va="bottom",
        color=annot.color,
        fontsize=9,
        fontweight="bold",
    )

ax.set_title("ROIs on false RGB preview")
plt.show()
# end region