# 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/classification#evaluate-training-roi-predictions

# region: setup
import os
from pathlib import Path

import joblib
import numpy as np
import qtec_hv_sdk as hs
import qtec_hv_sdk.annotations
from qtec_hv_sdk.preprocessing import make_reference
from qtec_hv_sdk.preprocessing import reflectance_calibration
from sklearn.metrics import confusion_matrix


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."
    )
TRAIN_CUBE = os.environ.get("HSI_EXAMPLE_TRAIN_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")
CLASSIFIER_MODEL_PATH = Path(os.environ.get("HSI_EXAMPLE_CLASSIFIER_MODEL", "pixel_classifier.joblib"))


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=TRAIN_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 annotation_value(value):
    return value[0] if value is not None else None


def load_annotations():
    annotations_path = Path(os.environ.get("HSI_EXAMPLE_ANNOTATIONS", "mix1_roi.json"))
    if not annotations_path.is_absolute():
        annotations_path = BASE_DIR / annotations_path
    if not annotations_path.exists():
        raise SystemExit("Set HSI_EXAMPLE_ANNOTATIONS to the annotations JSON file.")
    return hs.annotations.open(str(annotations_path))


def load_classifier_model():
    if not CLASSIFIER_MODEL_PATH.exists():
        raise SystemExit(
            f"Classifier model not found at {CLASSIFIER_MODEL_PATH}. "
            "Run 01_train_a_pixel_classifier_from_rois.py first, or set HSI_EXAMPLE_CLASSIFIER_MODEL."
        )
    return joblib.load(CLASSIFIER_MODEL_PATH)
# end region

# region: example
reflectance = open_reflectance_cube()
clf = load_classifier_model()
ann_file = load_annotations()

expected_list = []
predicted_list = []

for annot in ann_file.annotations:
    expected_class = annotation_value(annot.properties["type"])
    selected = reflectance.select_mask_from_descriptor(annot.descriptor)
    spectra = selected.to_numpy_with_interleave(hs.bip)[:, 0, :]

    expected_list.append(np.full(spectra.shape[0], expected_class))
    predicted_list.append(clf.predict(spectra))

expected = np.concatenate(expected_list)
predicted = np.concatenate(predicted_list)

classes = np.array(clf.classes_)
accuracy = np.mean(predicted == expected)
matrix = confusion_matrix(expected, predicted, labels=classes)
per_class_accuracy = matrix.diagonal() / np.maximum(matrix.sum(axis=1), 1)

print(f"Training ROI pixel accuracy: {accuracy:.3f}")
for class_name, class_accuracy in zip(classes, per_class_accuracy):
    print(f"{class_name}: {class_accuracy:.3f}")

print("Confusion matrix rows=true, columns=predicted")
print(classes)
print(matrix)
# end region