# 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/pca#pca-score-scatter-plot

# 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
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."
    )
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")
PCA_MODEL_PATH = Path(os.environ.get("HSI_EXAMPLE_PCA_MODEL", "pca_model.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 load_pca_model():
    if not PCA_MODEL_PATH.exists():
        raise SystemExit(
            f"PCA model not found at {PCA_MODEL_PATH}. "
            "Run 01_principal_component_analysis.py first, or set HSI_EXAMPLE_PCA_MODEL."
        )
    return joblib.load(PCA_MODEL_PATH)
# end region

# region: example
def sample_training_pixels(image, line_step=50, sample_step=20):
    sample_lines = []
    for line_index in range(0, min(200, image.shape.lines), line_step):
        line = image.array_plane(line_index, hs.lines).T
        sample_lines.append(line[::sample_step])
    return np.concatenate(sample_lines, axis=0)


reflectance = open_reflectance_cube()
sample_pixels = sample_training_pixels(reflectance)
pca = load_pca_model()

sample_scores = pca.transform(sample_pixels)

plt.figure(figsize=(8, 6))
plt.scatter(sample_scores[:, 0], sample_scores[:, 1], s=4, alpha=0.25)
plt.xlabel("PC1")
plt.ylabel("PC2")
plt.title("Sampled pixel scores")
plt.grid(True, alpha=0.3)
plt.show()
# end region