# 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#qamlib-capture-with-sdk-processing

# region: setup
import os
from pathlib import Path

import numpy as np
import qtec_hv_sdk as hs
from qtec_hv_sdk.preprocessing import make_reference
from qtec_hv_sdk.preprocessing import reflectance_calibration
from qtec_hv_sdk.util import predictor

import qamlib


BASE_DIR = Path(os.environ.get("HSI_EXAMPLE_BASE_DIR", "/path/to/HSI_data/nuts"))
N_LINES = int(os.environ.get("HSI_EXAMPLE_N_LINES", "100"))
VIDEO_DEVICE = os.environ.get("HSI_EXAMPLE_VIDEO_DEVICE")

EXPOSURE_US = os.environ.get("HSI_EXAMPLE_EXPOSURE_US", "1000")
FRAMERATE = os.environ.get("HSI_EXAMPLE_FRAMERATE")
QAMLIB_CROP = os.environ.get("HSI_EXAMPLE_QAMLIB_CROP")

DARK_REF = os.environ.get("HSI_EXAMPLE_DARK_REF")
WHITE_REF = os.environ.get("HSI_EXAMPLE_WHITE_REF")
CLASSIFIER_MODEL = os.environ.get("HSI_EXAMPLE_CLASSIFIER_MODEL")


def data_path(filename):
    path = Path(filename)
    if path.is_absolute():
        return path
    return BASE_DIR / path


def parse_crop(value):
    left, top, width, height = value.replace(":", ",").split(",", maxsplit=3)
    return int(left), int(top), int(width), int(height)


def open_qamlib_camera():
    cam = qamlib.Camera(VIDEO_DEVICE) if VIDEO_DEVICE else qamlib.Camera()
    cam.set_control("Exposure Time, Absolute", int(EXPOSURE_US))

    if FRAMERATE:
        cam.set_framerate(float(FRAMERATE))

    if QAMLIB_CROP:
        cam.set_crop(*parse_crop(QAMLIB_CROP))

    return cam


def normalize_qamlib_frame(frame):
    frame = np.asarray(frame)

    if frame.ndim == 3 and 1 in frame.shape:
        frame = np.squeeze(frame)
    if frame.ndim != 2:
        raise ValueError(f"Expected a 2D qamlib frame, got shape {frame.shape}")

    # The cropped qamlib frame is treated as one BIL camera line:
    # bands x samples. Add a line axis so the SDK sees L x B x S.
    return frame[np.newaxis, :, :]
# end region

# region: example
def make_line_processor():
    dark_ref = None
    white_ref = None
    if DARK_REF and WHITE_REF:
        dark_ref = make_reference(hs.open(str(data_path(DARK_REF)))).to_interleave(hs.bil).resolve()
        white_ref = make_reference(hs.open(str(data_path(WHITE_REF)))).to_interleave(hs.bil).resolve()

    hs_classifier = None
    if CLASSIFIER_MODEL:
        from joblib import load

        hs_classifier = predictor(load(data_path(CLASSIFIER_MODEL)))

    def process_frame(frame):
        line = normalize_qamlib_frame(frame)
        # qamlib returns a NumPy frame, not an SDK camera pipeline source.
        # Wrapping each line is the temporary bridge until direct local
        # capture can be used as an Image source inside the SDK.
        img = hs.Image.from_numpy(line, hs.bil)

        if dark_ref is not None and white_ref is not None:
            img = reflectance_calibration(img, white_ref, dark_ref)

        if hs_classifier is not None:
            img = hs_classifier(img)

        return img.to_numpy_with_interleave(hs.bip)[0]

    return process_frame


# This example is intended for on-camera execution. qamlib captures directly
# from the local V4L2 device, while the HV SDK performs HSI processing.
cam = open_qamlib_camera()
process_frame = make_line_processor()

with cam:
    for _ in range(N_LINES):
        meta, frame = cam.get_frame(buffered=True)

        if getattr(meta, "dropped", False):
            print(f"Dropped frame {meta.sequence}")
            continue

        processed_line = process_frame(frame)
        print(f"Frame {meta.sequence}: processed shape={processed_line.shape}")
# end region