# 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/streaming#camera-pipelines

# region: setup
import os
from pathlib import Path

import numpy as np
import qtec_hv_sdk as hs


BASE_DIR = Path(os.environ.get("HSI_EXAMPLE_BASE_DIR", "/path/to/HSI_data/nuts"))
TEST_CUBE = os.environ.get("HSI_EXAMPLE_TEST_CUBE", "mix2.pam")
CAMERA_IP = os.environ.get("HSI_EXAMPLE_CAMERA_IP")
N_LINES = int(os.environ.get("HSI_EXAMPLE_N_LINES", "10"))
EXPOSURE_US = os.environ.get("HSI_EXAMPLE_EXPOSURE_US")
FRAMERATE = os.environ.get("HSI_EXAMPLE_FRAMERATE")
HORIZONTAL_CROP = os.environ.get("HSI_EXAMPLE_HORIZONTAL_CROP")
BANDS = os.environ.get("HSI_EXAMPLE_BANDS")

CAMERA_WAVELENGTH_START_NM = float(
    os.environ.get("HSI_EXAMPLE_WAVELENGTH_START_NM", "430")
)
CAMERA_WAVELENGTH_END_NM = float(
    os.environ.get("HSI_EXAMPLE_WAVELENGTH_END_NM", "1700")
)

def add_camera_wavelengths(img):
    # Example PAM cubes may not include wavelengths; attach them so the
    # simulated camera exposes the same wavelength metadata as real hardware.
    meta = img.meta
    meta.spectral = hs.SpectralMeta(
        wavelengths=hs.WavelengthMeta(
            np.linspace(
                CAMERA_WAVELENGTH_START_NM,
                CAMERA_WAVELENGTH_END_NM,
                img.meta.shape.bands,
            ),
            hs.WavelengthUnit.Nanometer,
        )
    )
    return img.with_meta(meta)

def open_test_cube():
    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."
        )
    return add_camera_wavelengths(hs.open(str(BASE_DIR / TEST_CUBE)))


def parse_range(value):
    start, end = value.replace(":", ",").split(",", maxsplit=1)
    return int(start), int(end)
# end region

# region: example
def configure_camera(cam):
    if EXPOSURE_US:
        cam.set_exposure(int(EXPOSURE_US))
    if FRAMERATE:
        cam.set_framerate(float(FRAMERATE))
    if HORIZONTAL_CROP:
        start, end = parse_range(HORIZONTAL_CROP)
        cam.set_horizontal_crop(start, end)
    if BANDS:
        cam.set_bands([parse_range(part) for part in BANDS.split(";")])

# To use a real camera on the default ETH-B address:
# export HSI_EXAMPLE_CAMERA_IP=10.100.10.100
# When this variable is not set, the script uses a simulated camera backed by
# HSI_EXAMPLE_TEST_CUBE. That makes it possible to test the same pipeline
# without camera hardware.
def make_camera():
    if CAMERA_IP:
        cam = hs.control.Camera(CAMERA_IP)
        configure_camera(cam)
        return cam

    test_img = open_test_cube()
    return hs.control.Camera(test_img)


cam = make_camera()

wavelengths = np.array(cam.get_wavelengths())
band_650 = int(np.argmin(np.abs(wavelengths - 650.0)))
print(f"Closest band to 650 nm: {band_650}, wavelength={wavelengths[band_650]:.1f} nm")

# For real cameras, configure exposure, framerate, crop, and bands before
# creating the Image pipeline or starting a direct stream.
img = cam.to_hs_image()
processed = img.ensure_dtype(hs.float32) / 255.0

# The camera pipeline is executed when data is consumed.
# Slice the line dimension first; a live camera stream is otherwise open-ended.
first_lines = processed[:N_LINES, :, :].resolve()
first_band = first_lines[:, :, band_650]

print(f"{first_lines.shape=}")
print(f"{first_band.shape=}")

# end region