# 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-with-a-saved-regressor

# region: setup
import os
from pathlib import Path

import joblib
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


BASE_DIR = Path(os.environ.get("HSI_EXAMPLE_BASE_DIR", "/path/to/HSI_data/datasets"))
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."
    )
DARK_REF = os.environ.get("HSI_EXAMPLE_DARK_REF", "dark_ref.pam")
WHITE_REF = os.environ.get("HSI_EXAMPLE_WHITE_REF", "white_ref.pam")
MILK_TEST_CUBE = os.environ.get("HSI_EXAMPLE_MILK_TEST_CUBE", "milk.pam")
REGRESSION_MODEL_PATH = Path(os.environ.get("HSI_EXAMPLE_REGRESSION_MODEL", "regression_model.joblib"))
TARGET_PROPERTY = os.environ.get("HSI_EXAMPLE_TARGET_PROPERTY", "fat")

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 path_from_base(path):
    path = Path(path)
    if path.is_absolute():
        return path
    return BASE_DIR / path


def required_data_path(path, description):
    resolved = path_from_base(path)
    if not resolved.exists():
        raise SystemExit(
            f"Missing {description}: {resolved}\n"
            "This streaming regressor example uses the milk-fat dataset. Set "
            "HSI_EXAMPLE_BASE_DIR to the folder containing the milk cube and "
            "matching dark/white references."
        )
    return resolved


def make_references():
    dark = hs.open(str(required_data_path(DARK_REF, "dark reference")))
    white = hs.open(str(required_data_path(WHITE_REF, "white reference")))
    return make_reference(dark), make_reference(white)


def load_regression_model():
    if not REGRESSION_MODEL_PATH.exists():
        raise SystemExit(
            f"Regression model not found at {REGRESSION_MODEL_PATH}. "
            "Run the first regression example first, or set HSI_EXAMPLE_REGRESSION_MODEL."
        )
    return joblib.load(REGRESSION_MODEL_PATH)
# end region

# region: example
from qtec_hv_sdk.util import predictor


reg = load_regression_model()

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

# Simulated camera
test_img = add_camera_wavelengths(
    hs.open(str(required_data_path(MILK_TEST_CUBE, "milk datacube")))
)
cam = hs.control.Camera(test_img)

dark_ref, white_ref = make_references()

# Prepare calibration references for the camera pipeline.
# References interleave must match camera output.
dark_ref = dark_ref.to_interleave(hs.bil).resolve()
white_ref = white_ref.to_interleave(hs.bil).resolve()

# Create a lazy Image pipeline from the camera stream.
camera_img = cam.to_hs_image()

# Limit the stream to a finite number of frames for this example.
n_frames = 10
camera_img = camera_img[:n_frames, :, :]
camera_reflectance = reflectance_calibration(
    camera_img,
    white_ref,
    dark_ref,
    clip=True,
)
camera_reflectance = camera_reflectance.ensure_dtype(hs.float32).clip(1e-6, 1.0)
camera_absorbance = camera_reflectance.ufunc(lambda meta, plane: -np.log10(np.clip(plane, 1e-6, 1.0)))

hs_regressor = predictor(reg)
camera_prediction = hs_regressor(camera_absorbance)

with camera_prediction.stream() as stream:
    for meta, prediction_line in stream:
        values = prediction_line.ravel()
        print(
            f"Frame {meta.seq}: "
            f"{TARGET_PROPERTY} {values.min():.2f} to {values.max():.2f}"
        )
# end region