# 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#foreground-first-then-classify-objects

# region: setup
import os
from pathlib import Path

import matplotlib.pyplot as plt
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 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")
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")


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


class NumericLabelClassifier:
    # hs.util.predictor() returns an Image, so the prediction output must be
    # numeric image data. This adapter maps object class names to uint8 ids.
    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)
# end region

# region: train
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import LinearSVC


class NumericForegroundClassifier:
    def __init__(self, clf):
        self.clf = clf

    def predict(self, pixels):
        return self.clf.predict(pixels).astype(np.uint8)


reflectance = open_reflectance_cube()
ann_file = load_annotations()

pixels_list = []
labels_list = []

for annot in ann_file.annotations:
    class_name = annotation_value(annot.properties["type"])
    selected = reflectance.select_mask_from_descriptor(annot.descriptor)
    spectra = selected.to_numpy_with_interleave(hs.bip)[:, 0, :]
    pixels_list.append(spectra)
    labels_list.append(np.full(spectra.shape[0], class_name))

pixels = np.concatenate(pixels_list)
labels = np.concatenate(labels_list)

foreground_labels = labels != "background"

foreground_clf = make_pipeline(
    StandardScaler(),
    LogisticRegression(C=1.0, class_weight="balanced", max_iter=2000),
)
foreground_clf.fit(pixels, foreground_labels)

object_pixels = pixels[foreground_labels]
object_labels = labels[foreground_labels]

object_clf = make_pipeline(
    StandardScaler(),
    LinearSVC(C=1.0, class_weight="balanced", dual=False, max_iter=10000),
)
object_clf.fit(object_pixels, object_labels)
# end region

# region: test
test_reflectance = open_reflectance_cube(TEST_CUBE)

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

foreground_mask = predictor(NumericForegroundClassifier(foreground_clf))(preview_crop)
foreground_mask = foreground_mask.to_numpy_with_interleave(hs.bip)[:, :, 0].astype(bool)

numeric_object_clf = NumericLabelClassifier(object_clf)
object_label_map = predictor(numeric_object_clf)(preview_crop)
object_label_map = object_label_map.to_numpy_with_interleave(hs.bip)[:, :, 0].astype(np.uint8)

label_map = np.zeros_like(object_label_map, dtype=np.uint8)
label_map[foreground_mask] = object_label_map[foreground_mask] + 1
classes = np.array(["background", *numeric_object_clf.classes_])

plt.imshow(label_map, aspect="auto", interpolation="nearest")
plt.ylabel("Line")
plt.xlabel("Sample")
plt.title("Foreground first classification")
plt.colorbar(ticks=np.arange(len(classes)), label="Class id")
plt.show()
# end region