# 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-classifier

# 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
from qtec_hv_sdk.util import predictor


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")
DARK_REF = os.environ.get("HSI_EXAMPLE_DARK_REF", "dark_ref.pam")
WHITE_REF = os.environ.get("HSI_EXAMPLE_WHITE_REF", "white_ref.pam")
CLASSIFIER_MODEL_PATH = Path(os.environ.get("HSI_EXAMPLE_CLASSIFIER_MODEL", "pixel_classifier.joblib"))

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)))


def make_references():
    dark = hs.open(str(BASE_DIR / DARK_REF))
    white = hs.open(str(BASE_DIR / WHITE_REF))
    return make_reference(dark), make_reference(white)


def load_classifier_model():
    if not CLASSIFIER_MODEL_PATH.exists():
        raise SystemExit(
            "Set HSI_EXAMPLE_CLASSIFIER_MODEL to the saved classifier .joblib file."
        )
    return joblib.load(CLASSIFIER_MODEL_PATH)


class NumericLabelClassifier:
    # hs.util.predictor() returns an Image, so the prediction output must be
    # numeric image data. This adapter maps class names to uint8 ids.
    def __init__(self, clf):
        self.clf = clf
        self.classes_ = np.array(clf.classes_)
        self.class_to_id_ = {
            class_name: class_id
            for class_id, class_name in enumerate(self.classes_)
        }

    def predict(self, pixels):
        labels = self.clf.predict(pixels)
        return np.array([self.class_to_id_[label] for label in labels], dtype=np.uint8)
# end region

# region: example
numeric_clf = NumericLabelClassifier(load_classifier_model())

dark_ref, white_ref = make_references()

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

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

# Limit the stream to a finite number of frames for this example.
n_frames = 10
img = img[:n_frames, :, :]

# 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()
img = reflectance_calibration(img, white_ref, dark_ref, clip=True)

# Setup a predictor for the camera pipeline
pred = predictor(numeric_clf)
img = pred(img)

with img.stream() as stream:
    for meta, predicted in stream:
        predicted = predicted.astype(np.uint8).ravel()
        classes, counts = np.unique(predicted, return_counts=True)
        print(f"Frame {meta.seq}: {dict(zip(classes, counts.astype(int)))}")
# end region