# 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#lab-scanner-stage-controller

# 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")
STAGE_CONTROLLER_IP = os.environ.get("HSI_EXAMPLE_STAGE_CONTROLLER_IP")
N_LINES = int(os.environ.get("HSI_EXAMPLE_N_LINES", "10"))
FPS = float(os.environ.get("HSI_EXAMPLE_STAGE_FPS", "30"))
DISTANCE_MM = float(os.environ.get("HSI_EXAMPLE_STAGE_DISTANCE_MM", "50"))
OVERSAMPLING = float(os.environ.get("HSI_EXAMPLE_STAGE_OVERSAMPLING", "4"))
OUTPUT_PATH = Path(os.environ.get("HSI_EXAMPLE_OUTPUT", "lab_scanner_capture.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 scanner camera exposes the same wavelength metadata as hardware.
    if img.meta.spectral.wavelengths is not None:
        return img

    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 make_stage_controller():
    if STAGE_CONTROLLER_IP:
        return hs.control.StageController(STAGE_CONTROLLER_IP)

    test_img = open_test_cube()
    return hs.control.SimulatedStageController(
        test_img,
        fps=FPS,
        distance=DISTANCE_MM,
    )
# end region

# region: example
# To use real lab scanner hardware:
# export HSI_EXAMPLE_STAGE_CONTROLLER_IP=10.100.10.100
# When this variable is not set, the script uses a simulated stage controller
# backed by HSI_EXAMPLE_TEST_CUBE.
scanner = make_stage_controller()

scanner.set_fps(FPS)
scanner.distance = DISTANCE_MM
scanner.oversampling = OVERSAMPLING

wavelengths = np.array(scanner.hs_camera.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")

image = scanner.to_hs_image()
datacube = image[:N_LINES, :, :]
band_image = datacube[:, :, band_650]

# Consuming the lazy image starts the scanner workflow. With real hardware,
# hs.write(), resolve(), and to_numpy() all trigger capture. Use one terminal
# operation for the live scan; this example writes the scan to disk.
hs.write(datacube, str(OUTPUT_PATH))

# This would trigger a second capture
# array = band_image.to_numpy()

print(f"Saved lab scanner capture to {OUTPUT_PATH} with shape: {datacube.shape}")
# end region