# 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/regression#predict-a-regression-map

# 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")
DISPLAY_MIN = float(os.environ.get("HSI_EXAMPLE_REGRESSION_VMIN", "0.5"))
DISPLAY_MAX = float(os.environ.get("HSI_EXAMPLE_REGRESSION_VMAX", "36.0"))


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"
            "The regression examples use the milk-fat dataset. Set "
            "HSI_EXAMPLE_BASE_DIR to the folder containing the milk cube, "
            "milk_fat_roi.json, 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)


dark_ref, white_ref = make_references()


def open_absorbance_cube(cube_name):
    img = hs.open(str(required_data_path(cube_name, "milk datacube")))
    reflectance = reflectance_calibration(img, white_ref, dark_ref, clip=True)
    reflectance = reflectance.ensure_dtype(hs.float32).clip(1e-6, 1.0)
    return reflectance.ufunc(lambda meta, plane: -np.log10(np.clip(plane, 1e-6, 1.0)))


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

# region: example
import matplotlib.pyplot as plt
from qtec_hv_sdk.util import predictor


reg = load_regression_model()
test_absorbance = open_absorbance_cube(MILK_TEST_CUBE)

start_line = 0
n_lines = 800
test_crop = test_absorbance[start_line:start_line + n_lines, :, :]

hs_regressor = predictor(reg)
prediction = hs_regressor(test_crop)
prediction_map = prediction.to_numpy_with_interleave(hs.bip)[:, :, 0]
display_map = np.clip(prediction_map, DISPLAY_MIN, DISPLAY_MAX)

print(
    f"Raw predicted {TARGET_PROPERTY}: "
    f"{prediction_map.min():.2f} to {prediction_map.max():.2f}"
)

image = plt.imshow(
    display_map,
    cmap="viridis",
    vmin=DISPLAY_MIN,
    vmax=DISPLAY_MAX,
)
plt.colorbar(image, label=f"Predicted {TARGET_PROPERTY}", extend="both")
plt.title(f"{TARGET_PROPERTY} prediction, clipped for display")
plt.axis("off")
plt.show()
# end region