# 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#visualize-classification-results

# region: setup
import os
from pathlib import Path

import joblib
import matplotlib.pyplot as plt
import numpy as np
import qtec_hv_sdk as hs
import qtec_hv_sdk.annotations
from matplotlib.patches import Patch
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."
    )
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)


def hex_to_rgb(value):
    value = value.lstrip("#")
    return np.array(
        [int(value[0:2], 16), int(value[2:4], 16), int(value[4:6], 16)],
        dtype=np.uint8,
    )
# end region

# region: example
class NumericLabelClassifier:
    # hs.util.predictor() returns an Image, so the prediction output must be
    # numeric image data. This adapter keeps the original string-label
    # classifier, but maps class names to uint8 ids before returning labels.
    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)


ann_file = load_annotations()
clf = load_classifier_model()

train_reflectance = open_reflectance_cube()

class_colors_rgb = {
    annotation_value(annot.properties["type"]): hex_to_rgb(annot.color)
    for annot in ann_file.annotations
}

start_line = 0
n_lines = 300
overlay_alpha = 0.5

stop_line = min(start_line + n_lines, train_reflectance.shape.lines)
preview_crop = train_reflectance[start_line:stop_line, :, :]

numeric_clf = NumericLabelClassifier(clf)
classified = predictor(numeric_clf)(preview_crop)
label_map = classified.to_numpy_with_interleave(hs.bip)[:, :, 0].astype(np.uint8)

color_map = np.zeros((len(numeric_clf.classes_), 3), dtype=np.uint8)
for class_id, class_name in enumerate(numeric_clf.classes_):
    color_map[class_id] = class_colors_rgb[class_name]

class_rgb = color_map[label_map]

gray = preview_crop.mean_axis(hs.bands).to_numpy_with_interleave(hs.bip)
gray = np.clip(gray[:, :, 0] * 255, 0, 255).astype(np.uint8)
gray_rgb = np.stack([gray, gray, gray], axis=-1)

preview = (overlay_alpha * class_rgb + (1 - overlay_alpha) * gray_rgb).astype(np.uint8)

legend_handles = [
    Patch(
        color=color_map[class_id] / 255.0,
        label=str(class_name),
    )
    for class_id, class_name in enumerate(numeric_clf.classes_)
]

fig, ax = plt.subplots()
ax.imshow(preview)
ax.set_title("Training cube classification overlay")
ax.axis("off")
ax.legend(
    handles=legend_handles,
    loc="upper left",
    bbox_to_anchor=(1.02, 1.0),
    borderaxespad=0,
    title="Class",
)
fig.tight_layout()
plt.show()
# end region