# 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/basics#false-rgb-preview

# region: setup
import os
from pathlib import Path

import matplotlib.pyplot as plt
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/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."
    )
CUBE = os.environ.get("HSI_EXAMPLE_CUBE", "mix1.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")

WAVELENGTH_MIN_NM = 430.0
WAVELENGTH_MAX_NM = 1700.0


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 open_reflectance_cube(cube_name=CUBE):
    dark_ref, white_ref = make_references()
    img = hs.open(str(BASE_DIR / cube_name))
    return reflectance_calibration(img, white_ref, dark_ref, clip=True)


def wavelengths_for_bands(n_bands):
    # This example shows the manual band-wavelength mapping. In normal SDK
    # workflows, prefer wavelength metadata from img.meta.spectral.wavelengths.
    return np.linspace(WAVELENGTH_MIN_NM, WAVELENGTH_MAX_NM, n_bands)


def band_index_for_wavelength(wavelength_nm, n_bands):
    wavelengths = wavelengths_for_bands(n_bands)
    return int(np.argmin(np.abs(wavelengths - wavelength_nm)))


def contrast_stretch(image, percentiles=(1, 99)):
    low, high = np.percentile(image, percentiles)
    return np.clip((image - low) / (high - low + 1e-8), 0, 1)
# end region

# region: example
reflectance = open_reflectance_cube()

target_wavelengths = (650, 550, 460)
band_indices = [
    band_index_for_wavelength(wv, reflectance.shape.bands)
    for wv in target_wavelengths
]

sorted_band_indices = sorted(band_indices)
selected_cube = reflectance.select_bands(sorted_band_indices).to_numpy_with_interleave(hs.bip)

rgb = np.empty((*selected_cube.shape[:2], 3), dtype=np.float32)
for channel, band_index in enumerate(band_indices):
    selected_channel = sorted_band_indices.index(band_index)
    rgb[:, :, channel] = contrast_stretch(selected_cube[:, :, selected_channel])

plt.imshow(rgb)
plt.title(f"False RGB bands: {band_indices}")
plt.axis("off")
plt.show()
# end region