# 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#stream-individual-frames

# 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"))
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."
    )
TEST_CUBE = os.environ.get("HSI_EXAMPLE_TEST_CUBE", "mix2.pam")

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():
    return add_camera_wavelengths(hs.open(str(BASE_DIR / TEST_CUBE)))
# end region

# region: example
# Real camera
# cam = hs.control.Camera("10.100.10.100")

# Simulated camera
test_img = open_test_cube()
cam = hs.control.Camera(test_img)

with cam.stream(n_frames=10) as stream:
    while True:
        item = stream.get_frame()
        if item is None:
            break

        meta, frame = item

        if meta.dropped:
            print(f"Dropped frame {meta.seq} at {meta.timestamp}")
            continue

        print(f"Frame {meta.seq}: shape={frame.shape}")
# end region