# 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#classify-with-spectral-angle-mapper

# region: setup
import os
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import qtec_hv_sdk as hs
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")
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")
MEAN_SPECTRA_CSV = os.environ.get("HSI_EXAMPLE_MEAN_SPECTRA", "mix1_spectra.csv")
CLASS_PROPERTY = os.environ.get("HSI_EXAMPLE_CLASS_PROPERTY", "type")

WAVELENGTH_MIN_NM = 430.0
WAVELENGTH_MAX_NM = 1700.0


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 wavelengths_for_bands(n_bands):
    return np.linspace(WAVELENGTH_MIN_NM, WAVELENGTH_MAX_NM, n_bands)


def path_from_base(path):
    path = Path(path)
    if path.is_absolute():
        return path
    return BASE_DIR / path


def load_mean_spectra(csv_path, n_bands):
    csv_path = Path(csv_path) if Path(csv_path).is_absolute() else BASE_DIR / csv_path
    if not csv_path.exists():
        raise SystemExit("Set HSI_EXAMPLE_MEAN_SPECTRA to the exported mean spectra CSV file.")

    table = pd.read_csv(csv_path)
    wavelength_columns = sorted(
        [c for c in table.columns if c.replace(".", "", 1).isdigit()], key=float
    )
    if "name" not in table.columns or not wavelength_columns:
        raise SystemExit(
            "Mean spectra CSV must have a 'name' column and numeric wavelength columns."
        )

    spectra = table[["name", *wavelength_columns]].groupby("name", sort=True).mean()
    roi_names = spectra.index.to_numpy()
    csv_wavelengths = np.array([float(c) for c in wavelength_columns])
    cube_wavelengths = wavelengths_for_bands(n_bands)
    values = np.vstack([
        np.interp(cube_wavelengths, csv_wavelengths, row)
        for row in spectra.to_numpy()
    ]).astype(np.float32)
    return roi_names, values

# end region

# region: example
class SpectralAngleMapper:
    def __init__(self, reference_spectra, reference_labels):
        self.classes_ = np.unique(reference_labels)
        self.class_to_id_ = {
            class_name: class_id
            for class_id, class_name in enumerate(self.classes_)
        }
        self.reference_labels_ = np.array(reference_labels)
        self.reference_class_ids_ = np.array([
            self.class_to_id_[label]
            for label in self.reference_labels_
        ], dtype=np.uint8)

        reference_spectra = np.asarray(reference_spectra, dtype=np.float32)
        reference_norm = np.linalg.norm(reference_spectra, axis=1, keepdims=True)
        self.reference_spectra_ = reference_spectra / np.maximum(reference_norm, 1e-12)

    def predict(self, pixels):
        pixels = np.asarray(pixels, dtype=np.float32)
        pixel_norm = np.linalg.norm(pixels, axis=1, keepdims=True)
        normalized_pixels = pixels / np.maximum(pixel_norm, 1e-12)

        # Maximizing cosine similarity is equivalent to minimizing spectral angle.
        cosine = normalized_pixels @ self.reference_spectra_.T
        best_reference = np.argmax(cosine, axis=1)
        return self.reference_class_ids_[best_reference]


test_reflectance = open_reflectance_cube(TEST_CUBE)

roi_names, mean_spectra = load_mean_spectra(
    MEAN_SPECTRA_CSV,
    n_bands=test_reflectance.shape.bands,
)

table = pd.read_csv(BASE_DIR / MEAN_SPECTRA_CSV)
if CLASS_PROPERTY not in table.columns:
    raise SystemExit(
        f"The exported mean spectra file must include a '{CLASS_PROPERTY}' column. "
        "Export spectra from HV Explorer with properties included."
    )

class_by_name = dict(zip(table["name"], table[CLASS_PROPERTY]))
reference_labels = np.array([class_by_name[name] for name in roi_names])

sam = SpectralAngleMapper(mean_spectra, reference_labels)

start_line = 0
n_lines = 300
overlay_alpha = 0.5

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

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

color_map = np.zeros((len(sam.classes_), 3), dtype=np.uint8)
cmap = plt.get_cmap("tab10", len(sam.classes_))
for class_id, _class_name in enumerate(sam.classes_):
    color_map[class_id] = (np.array(cmap(class_id)[:3]) * 255).astype(np.uint8)

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(sam.classes_)
]

fig, ax = plt.subplots()
ax.imshow(preview)
ax.set_title("Spectral angle mapper from mean spectra")
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