# 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#mean-spectrum-from-an-roi

# 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


WAVELENGTH_MIN_NM = 430.0
WAVELENGTH_MAX_NM = 1700.0

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")
ROI_TYPE = os.environ.get("HSI_EXAMPLE_ROI_TYPE", "almond")


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)
# end region

# region: example
def annotation_value(value):
    return value[0] if value is not None else None


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 find_roi(ann_file, cube_name=CUBE, roi_type=ROI_TYPE):
    for annot in ann_file.annotations:
        image_path = Path(annot.image_path or cube_name).name
        if image_path != Path(cube_name).name:
            continue
        if annotation_value(annot.properties.get("type")) == roi_type:
            return annot
    raise SystemExit(f"No ROI with type {roi_type!r} found for {cube_name!r}.")


reflectance = open_reflectance_cube()
ann_file = load_annotations()
roi = find_roi(ann_file)
selected = reflectance.select_mask_from_descriptor(roi.descriptor)
spectra = selected.to_numpy_with_interleave(hs.bip)[:, 0, :]
mean = spectra.mean(axis=0)
std = spectra.std(axis=0)
wavelengths = wavelengths_for_image(reflectance)

plt.plot(wavelengths, mean, label="mean")
plt.fill_between(wavelengths, mean - std, mean + std, alpha=0.25, label="+/- std")
plt.xlabel("Wavelength [nm]")
plt.ylabel("Reflectance")
plt.title(f"ROI mean spectrum: {roi.title}")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()
# end region