# 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#load-a-model-exported-from-hv-explorer

# region: setup
import os
from pathlib import Path

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


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."
    )
TEST_CUBE = os.environ.get("HSI_EXAMPLE_TEST_CUBE", "mix2.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")
HV_EXPLORER_MODEL = os.environ.get("HSI_EXAMPLE_HV_EXPLORER_MODEL", "mix1_svc.pkl")
HV_EXPLORER_ANNOTATIONS = os.environ.get("HSI_EXAMPLE_HV_EXPLORER_ANNOTATIONS", "mix1_roi.json")


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=TEST_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 path_from_base(path):
    path = Path(path)
    if path.is_absolute():
        return path
    return BASE_DIR / path


def load_hv_explorer_annotations():
    annotations_path = path_from_base(HV_EXPLORER_ANNOTATIONS)
    if not annotations_path.exists():
        raise SystemExit(
            f"Annotations file not found at {annotations_path}. "
            "Export the annotations JSON from HV Explorer and set HSI_EXAMPLE_HV_EXPLORER_ANNOTATIONS."
        )
    return hs.annotations.open(str(annotations_path))
# end region

# region: example
import pickle

from qtec_hv_sdk.util import predictor


def load_hv_explorer_model():
    model_path = path_from_base(HV_EXPLORER_MODEL)
    if not model_path.exists():
        raise SystemExit(
            f"HV Explorer model not found at {model_path}. "
            "Export a classifier from HV Explorer and set HSI_EXAMPLE_HV_EXPLORER_MODEL."
        )
    with open(model_path, "rb") as f:
        return pickle.load(f)


hv_model = load_hv_explorer_model()

# HV Explorer stores the fitted sklearn-compatible model together with the
# annotation property that was used as the class label.
clf = hv_model["model"]
property_name = hv_model["property_name"]

ann_file = load_hv_explorer_annotations()

# The exported model predicts numeric class IDs. The annotation file keeps the
# human-readable label names in the same order.
labels = ann_file.property_desc[property_name].labels
id_to_label = dict(enumerate(labels))

print(f"Model: {hv_model['name']!r}  type: {hv_model['type']}")
print(f"Property: {property_name!r}")
print(f"Classes: {id_to_label}")

start_line = 0
n_lines = 300

test_reflectance = open_reflectance_cube()
crop = test_reflectance[start_line:start_line + n_lines, :, :]

# predictor() lets the model run directly on the lazy Image crop.
hs_clf = predictor(clf)
prediction = hs_clf(crop)
label_map = prediction.to_numpy_with_interleave(hs.bip)[:, :, 0].astype(int)

ids, counts = np.unique(label_map, return_counts=True)
print("\nPer-class pixel counts:")
for class_id, count in zip(ids, counts):
    label = id_to_label.get(int(class_id), str(class_id))
    print(f"  {label}: {count}")
# end region