# 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/basics#anomaly-detection-from-background-pixels

# 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
from sklearn.covariance import EmpiricalCovariance


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."
    )
CUBE = os.environ.get("HSI_EXAMPLE_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")
ANNOTATIONS = os.environ.get("HSI_EXAMPLE_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=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 contrast_stretch(image, percentiles=(1, 99)):
    low, high = np.percentile(image, percentiles)
    return np.clip((image - low) / (high - low + 1e-8), 0, 1)


def annotation_value(value):
    return value[0] if value is not None else None


def load_annotations():
    annotations_path = Path(ANNOTATIONS)
    if not annotations_path.is_absolute():
        annotations_path = BASE_DIR / annotations_path
    if not annotations_path.exists():
        raise SystemExit("Set HSI_EXAMPLE_ANNOTATIONS to an annotations JSON file with a background ROI.")
    return hs.annotations.open(str(annotations_path))


def find_background_annotation(ann_file):
    for annot in ann_file.annotations:
        if annotation_value(annot.properties.get("type")) == "background":
            return annot
    raise SystemExit("No background ROI found in the annotations file.")
# end region


# region: example
class MahalanobisDistance:
    def __init__(self, covariance_model):
        self.covariance_model = covariance_model

    def predict(self, pixels):
        return self.covariance_model.mahalanobis(pixels).reshape(-1, 1)


reflectance = open_reflectance_cube()

ann_file = load_annotations()
background_annotation = find_background_annotation(ann_file)
background = reflectance.select_mask_from_descriptor(background_annotation.descriptor)
# Mask selection returns a compact Image containing only the selected pixels.
# In BIP order this is [selected_pixels, 1, bands].
background_pixels = background.to_numpy_with_interleave(hs.bip)[:, 0, :]

band_step = 8
band_indices = list(range(0, background_pixels.shape[1], band_step))
model = EmpiricalCovariance()
model.fit(background_pixels[:, band_indices])

preview_cube = reflectance[0:250, 0:400, :].select_bands(band_indices)
# Wrap the sklearn-style model so it can run lazily over an Image pipeline.
distance_pipeline = predictor(MahalanobisDistance(model))(preview_cube)

distance = distance_pipeline.to_numpy_with_interleave(hs.bip)[:, :, 0]

plt.imshow(contrast_stretch(distance), cmap="magma")
plt.title("Anomaly map")
plt.axis("off")
plt.show()
# end region