rss_core 0.5.0

//! Pixel quality mask decoding helpers.
//!
//! Provides platform-independent functions to decode per-pixel QA/classification
//! bands (Sentinel-2 SCL, Landsat QA_PIXEL) into boolean masks, and apply those
//! masks to computed raster values.
//!
//! # Usage
//!
//! Inside an eorst worker function:
//!
//! ```rust,ignore
//! use ndarray::{s, Array4, Array3, Axis};
//! use rss_core::cloud_mask::{decode_qa_mask, apply_mask, QaType, MaskClasses};
//!
//! fn ndvi_worker(data: &Array4<i16>, dim: eorst::Dimension) -> Array3<i16> {
//!     let (t, _, r, c) = data.data().dim();
//!     let mut result = Array3::zeros((t, r, c));
//!
//!     for ti in 0..t {
//!         let red = data.slice(s![ti, 0, .., ..]).mapv(|v| v as f32);
//!         let nir = data.slice(s![ti, 1, .., ..]).mapv(|v| v as f32);
//!         let mut ndvi = ((nir - red) / (nir + red)).mapv(|v| (v * 10000.) as i16);
//!
//!         // QA band is at layer index 2
//!         let qa = data.slice(s![ti, 2, .., ..]);
//!         let mask = decode_qa_mask(qa, QaType::LandsatPixelQA, MaskClasses::default());
//!         apply_mask(ndvi.view_mut(), mask.view(), 0i16);
//!
//!         result.index_mut(Axis(0), ti).assign(&ndvi);
//!     }
//!     result
//! }
//! ```

use bitflags::bitflags;
use ndarray::{Array2, ArrayView2, ArrayViewMut2, Zip};

bitflags! {
    /// Bitmask controlling which pixel classes are masked out.
    ///
    /// Default = all flags set (cloud + shadow + water + snow).
    /// Clear individual flags to keep those pixels.
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct MaskClasses: u8 {
        /// Cloud pixels (all confidence levels)
        const CLOUD        = 0b0000_0001;
        /// Cloud shadow pixels
        const CLOUD_SHADOW = 0b0000_0010;
        /// Water bodies
        const WATER        = 0b0000_0100;
        /// Snow / ice cover
        const SNOW         = 0b0000_1000;
        /// All common non-vegetation classes
        const ALL          = Self::CLOUD.bits() | Self::CLOUD_SHADOW.bits()
                          | Self::WATER.bits() | Self::SNOW.bits();
    }
}

impl Default for MaskClasses {
    fn default() -> Self {
        Self::ALL
    }
}

/// Identifies the type of QA/classification band being decoded.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QaType {
    /// Sentinel-2 Scene Classification Layer (SCL)
    ///
    /// Single-byte values: 0-12 (plus 254/255 for NA).
    /// See MSI Level-2A Product User Manual for the full class mapping.
    S2SCL,

    /// Landsat 8/9 QA_PIXEL band
    ///
    /// 16-bit value with 2-bit flags packed per class:
    /// fill, cloud shadow, water, cloud, cloud conf, snow, clear, auto-flags.
    LandsatPixelQA,
}

/// Decode a 2D QA/classification slice into a boolean mask.
///
/// Returns an `Array2<bool>` of the same shape as `qa`, where `true` means
/// the pixel should be **masked out** (excluded from analysis).
///
/// # Arguments
///
/// * `qa` — 2D view of the QA band data for one time slice
/// * `qa_type` — which decoding logic to apply (S2 SCL or Landsat QA_PIXEL)
/// * `classes` — which pixel classes to mask; default is all
///
/// # Example
///
/// ```rust,ignore
/// let qa = data.slice(s![ti, qa_layer, .., ..]);
/// let mask = decode_qa_mask(qa, QaType::S2SCL, MaskClasses::ALL);
/// ```
pub fn decode_qa_mask(
    qa: ArrayView2<i16>,
    qa_type: QaType,
    classes: MaskClasses,
) -> Array2<bool> {
    match qa_type {
        QaType::S2SCL => qa.mapv(|v| s2_scl_is_masked(v, classes)),
        QaType::LandsatPixelQA => qa.mapv(|v| landsat_qa_is_masked(v, classes)),
    }
}

/// Apply a boolean mask to a result array.
///
/// Where `mask` is `true`, the corresponding element in `result` is set to `na`.
///
/// # Arguments
///
/// * `result` — mutable 2D view of the output data
/// * `mask` — boolean mask (same shape as `result`), `true` = masked out
/// * `na` — no-data value to write for masked pixels
pub fn apply_mask<T: Copy>(result: ArrayViewMut2<T>, mask: ArrayView2<bool>, na: T) {
    Zip::from(result)
        .and(mask)
        .for_each(|r, m| {
            if *m {
                *r = na;
            }
        });
}

// ---------------------------------------------------------------------------
// Per-pixel decoders (internal)
// ---------------------------------------------------------------------------

/// Returns `true` if the SCL pixel should be masked.
fn s2_scl_is_masked(value: i16, classes: MaskClasses) -> bool {
    match value {
        // Always mask: no data, saturation, undefined
        0 | 1 | 254 | 255 => true,
        // Cloud shadow (class 3)
        3 => classes.contains(MaskClasses::CLOUD_SHADOW),
        // Vegetation (4), Not vegetation (5) — always clear
        4 | 5 => false,
        // Water (class 6)
        6 => classes.contains(MaskClasses::WATER),
        // Unclassified (7) — keep
        7 => false,
        // Cloud low probability (8) — keep (not masked by default)
        8 => false,
        // Cloud medium (9), high (10), thin cloud (11) probability
        9 | 10 | 11 => classes.contains(MaskClasses::CLOUD),
        // Snow (class 12)
        12 => classes.contains(MaskClasses::SNOW),
        // Anything else — keep
        _ => false,
    }
}

/// Returns `true` if the Landsat QA_PIXEL value should be masked.
///
/// Bit layout (2 bits per flag, little-endian):
/// | Bits | Field |
/// |------|-------|
/// | 0-1  | Fill |
/// | 2-3  | Cloud Shadow |
/// | 4-5  | Water |
/// | 6-7  | Cloud |
/// | 8-9  | Cloud Confidence |
/// | 10-11| Snow |
/// | 12-13| Clear |
/// | 14-15| Auto-Flags |
///
/// Each 2-bit flag: 0 = not defined, 1 = not present, 2 = present/low, 3 = high
fn landsat_qa_is_masked(value: i16, classes: MaskClasses) -> bool {
    if value < 0 {
        return true; // fill / invalid
    }

    let cloud_shadow = (value >> 2) & 0x3;
    let water = (value >> 4) & 0x3;
    let cloud = (value >> 6) & 0x3;
    let cloud_conf = (value >> 8) & 0x3;
    let snow = (value >> 10) & 0x3;

    let is_cloud = cloud >= 2 || cloud_conf >= 2;
    let is_shadow = cloud_shadow >= 2;
    let is_water = water >= 2;
    let is_snow = snow >= 2;

    (classes.contains(MaskClasses::CLOUD) && is_cloud)
        || (classes.contains(MaskClasses::CLOUD_SHADOW) && is_shadow)
        || (classes.contains(MaskClasses::WATER) && is_water)
        || (classes.contains(MaskClasses::SNOW) && is_snow)
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::Array2;

    // ---- S2 SCL tests ----

    #[test]
    fn s2_scl_clear_vegetation_not_masked() {
        let qa = Array2::from_shape_vec((2, 2), vec![4, 5, 4, 5]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::S2SCL, MaskClasses::ALL);
        assert!(mask.iter().all(|m| !m));
    }

    #[test]
    fn s2_scl_cloud_masked() {
        let qa = Array2::from_shape_vec((1, 3), vec![9, 10, 11]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::S2SCL, MaskClasses::ALL);
        assert!(mask.iter().all(|m| *m));
    }

    #[test]
    fn s2_scl_shadow_masked() {
        let qa = Array2::from_shape_vec((1, 1), vec![3]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::S2SCL, MaskClasses::ALL);
        assert!(mask.iter().all(|m| *m));
    }

    #[test]
    fn s2_scl_water_masked() {
        let qa = Array2::from_shape_vec((1, 1), vec![6]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::S2SCL, MaskClasses::ALL);
        assert!(mask.iter().all(|m| *m));
    }

    #[test]
    fn s2_scl_snow_masked() {
        let qa = Array2::from_shape_vec((1, 1), vec![12]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::S2SCL, MaskClasses::ALL);
        assert!(mask.iter().all(|m| *m));
    }

    #[test]
    fn s2_scl_selective_cloud_only() {
        // With only CLOUD flag: shadow(3) and water(6) should NOT be masked
        let qa = Array2::from_shape_vec((1, 4), vec![3, 6, 10, 12]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::S2SCL, MaskClasses::CLOUD);
        assert!(!mask[[0, 0]]); // shadow — not masked
        assert!(!mask[[0, 1]]); // water — not masked
        assert!(mask[[0, 2]]);  // cloud — masked
        assert!(!mask[[0, 3]]); // snow — not masked
    }

    #[test]
    fn s2_scl_nodata_always_masked() {
        let qa = Array2::from_shape_vec((1, 4), vec![0, 1, 254, 255]).unwrap();
        // Even with empty mask classes, no-data is always masked
        let mask = decode_qa_mask(qa.view(), QaType::S2SCL, MaskClasses::empty());
        assert!(mask.iter().all(|m| *m));
    }

    // ---- Landsat QA_PIXEL tests ----

    #[test]
    fn landsat_qa_clear_not_masked() {
        // Clear = bits 12-13 set to 2: 0b0010 << 12 = 8192
        // 2-bit flags: 0=not defined, 1=not present, 2=low/present, 3=high
        let qa = Array2::from_shape_vec((1, 1), vec![8192]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::LandsatPixelQA, MaskClasses::ALL);
        assert!(!mask[[0, 0]]);
    }

    #[test]
    fn landsat_qa_cloud_masked() {
        // Cloud present (low conf) = bits 6-7 set to 2: 0b0010 << 6 = 128
        let qa = Array2::from_shape_vec((1, 1), vec![128]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::LandsatPixelQA, MaskClasses::ALL);
        assert!(mask[[0, 0]]);
    }

    #[test]
    fn landsat_qa_shadow_masked() {
        // Cloud shadow present = bits 2-3 set to 2: 0b0010 << 2 = 8
        let qa = Array2::from_shape_vec((1, 1), vec![8]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::LandsatPixelQA, MaskClasses::ALL);
        assert!(mask[[0, 0]]);
    }

    #[test]
    fn landsat_qa_water_masked() {
        // Water present = bits 4-5 set to 2: 0b0010 << 4 = 32
        let qa = Array2::from_shape_vec((1, 1), vec![32]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::LandsatPixelQA, MaskClasses::ALL);
        assert!(mask[[0, 0]]);
    }

    #[test]
    fn landsat_qa_snow_masked() {
        // Snow present = bits 10-11 set to 2: 0b0010 << 10 = 2048
        let qa = Array2::from_shape_vec((1, 1), vec![2048]).unwrap();
        let mask = decode_qa_mask(qa.view(), QaType::LandsatPixelQA, MaskClasses::ALL);
        assert!(mask[[0, 0]]);
    }

    #[test]
    fn landsat_qa_selective_water_only() {
        // With only WATER flag: cloud(128) and shadow(8) should NOT be masked
        let qa = Array2::from_shape_vec((1, 3), vec![8, 32, 128]).unwrap(); // shadow, water, cloud
        let mask = decode_qa_mask(qa.view(), QaType::LandsatPixelQA, MaskClasses::WATER);
        assert!(!mask[[0, 0]]); // shadow — not masked
        assert!(mask[[0, 1]]);  // water — masked
        assert!(!mask[[0, 2]]); // cloud — not masked
    }

    // ---- apply_mask test ----

    #[test]
    fn apply_mask_sets_na_values() {
        let mut data = Array2::from_shape_vec((2, 2), vec![100, 200, 300, 400]).unwrap();
        let mask = Array2::from_shape_vec((2, 2), vec![true, false, false, true]).unwrap();
        apply_mask(data.view_mut(), mask.view(), -9999i16);
        assert_eq!(data[[0, 0]], -9999);
        assert_eq!(data[[0, 1]], 200);
        assert_eq!(data[[1, 0]], 300);
        assert_eq!(data[[1, 1]], -9999);
    }
}