# 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#spatial-and-spectral-binning

# region: setup
import os
from pathlib import Path

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


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)
# end region

# region: example
reflectance = open_reflectance_cube()

# Average groups of 4 neighboring samples (spatial columns)
sample_binned = reflectance.binning(4, hs.samples)

# Average groups of 4 neighboring lines (spatial rows)
line_binned = reflectance.binning(4, hs.lines)

# Average groups of 4 neighboring lines and samples (full 4x4 spatial binning)
spatial_binned = reflectance.binning(4, hs.lines).binning(4, hs.samples)

# Average groups of 6 neighboring bands (spectral channels)
spectral_binned = reflectance.binning(6, hs.bands)

print(sample_binned.shape)
print(line_binned.shape)
print(spatial_binned.shape)
print(spectral_binned.shape)
# end region