neurodoom 0.6.7

//! AI-facing snapshot of the engine's semantic + depth buffers.
//!
//! A [`PerceptionFrame`] is the minimum information a learned controller
//! needs to act: which class each pixel belongs to (wall, enemy, door,
//! ...) and how far it is. Built from a rendered frame via
//! [`PerceptionFrame::from_engine`], optionally downsampled with
//! [`PerceptionFrame::downsample`], and queried with helpers like
//! [`PerceptionFrame::columns_with_class`] and
//! [`PerceptionFrame::nearest_depth_of`].

use alloc::vec;
use alloc::vec::Vec;

use crate::engine::DoomEngine;
use crate::math::{Fixed, SCREENHEIGHT, SCREENWIDTH};
pub use crate::render::SemanticClass;

/// One frame of AI-readable perception data.
pub struct PerceptionFrame {
    pub tick: u32,
    pub width: usize,
    pub height: usize,
    /// SemanticClass per pixel (row-major, width × height).
    pub semantic: Vec<u8>,
    /// Distance per pixel in fixed-point map units (row-major).
    pub depth: Vec<Fixed>,
}

impl PerceptionFrame {
    /// Capture a perception frame from the current engine state.
    pub fn from_engine<R: crate::rules::GameRules>(engine: &DoomEngine<R>) -> Self {
        Self {
            tick: engine.world.tick,
            width: SCREENWIDTH,
            height: SCREENHEIGHT,
            semantic: engine.semantic_buffer().to_vec(),
            depth: engine.depth_buffer().to_vec(),
        }
    }

    /// Downsample to a smaller resolution.
    /// Semantic: majority vote per block. Depth: mean per block.
    pub fn downsample(&self, new_w: usize, new_h: usize) -> Self {
        if new_w == 0 || new_h == 0 || new_w > self.width || new_h > self.height {
            return Self {
                tick: self.tick,
                width: new_w,
                height: new_h,
                semantic: vec![SemanticClass::Void as u8; new_w * new_h],
                depth: vec![0; new_w * new_h],
            };
        }

        let bw = self.width / new_w;
        let bh = self.height / new_h;
        let mut semantic = vec![0u8; new_w * new_h];
        let mut depth = vec![0i32; new_w * new_h];

        for ny in 0..new_h {
            for nx in 0..new_w {
                let out_idx = ny * new_w + nx;

                // Collect classes and depths in this block
                let mut counts = [0u32; 9]; // one per SemanticClass
                let mut depth_sum: i64 = 0;
                let mut depth_count = 0u32;

                for by in 0..bh {
                    for bx in 0..bw {
                        let sy = ny * bh + by;
                        let sx = nx * bw + bx;
                        if sy < self.height && sx < self.width {
                            let src_idx = sy * self.width + sx;
                            if let Some(&cls) = self.semantic.get(src_idx)
                                && let Some(slot) = counts.get_mut(cls as usize)
                            {
                                *slot = slot.saturating_add(1);
                            }
                            if let Some(&d) = self.depth.get(src_idx)
                                && d != 0
                            {
                                depth_sum = depth_sum.saturating_add(d as i64);
                                depth_count = depth_count.saturating_add(1);
                            }
                        }
                    }
                }

                // Majority vote for semantic class
                let best_class = counts
                    .iter()
                    .enumerate()
                    .max_by_key(|(_, c)| **c)
                    .map(|(i, _)| i as u8)
                    .unwrap_or(0);
                if let Some(s) = semantic.get_mut(out_idx) {
                    *s = best_class;
                }

                // Mean depth
                if let Some(slot) = depth.get_mut(out_idx) {
                    *slot = if depth_count > 0 {
                        (depth_sum / depth_count as i64) as Fixed
                    } else {
                        0
                    };
                }
            }
        }

        Self {
            tick: self.tick,
            width: new_w,
            height: new_h,
            semantic,
            depth,
        }
    }

    /// Returns screen columns (x indices) that contain a given class.
    pub fn columns_with_class(&self, class: SemanticClass) -> Vec<usize> {
        let target = class as u8;
        // Iterate row-by-row, then pull out the column by position.
        let mut seen = alloc::vec![false; self.width];
        for row in self.semantic.chunks_exact(self.width) {
            for (x, &cls) in row.iter().enumerate() {
                if cls == target
                    && let Some(flag) = seen.get_mut(x)
                {
                    *flag = true;
                }
            }
        }
        seen.iter()
            .enumerate()
            .filter_map(|(x, &s)| s.then_some(x))
            .collect()
    }

    /// Nearest depth of a given class across the whole frame.
    pub fn nearest_depth_of(&self, class: SemanticClass) -> Option<Fixed> {
        let target = class as u8;
        self.semantic
            .iter()
            .zip(self.depth.iter())
            .filter_map(|(&cls, &d)| (cls == target && d > 0).then_some(d))
            .min()
    }

    /// Count pixels of a given class.
    pub fn count_class(&self, class: SemanticClass) -> usize {
        let target = class as u8;
        self.semantic.iter().filter(|&&c| c == target).count()
    }
}

#[cfg(test)]
// Tests deliberately use direct indexing and literal multipliers for
// readability; the lib-level lints are aimed at production code.
#[allow(clippy::indexing_slicing, clippy::erasing_op, clippy::identity_op)]
mod tests {
    use super::*;

    fn make_frame() -> PerceptionFrame {
        let w = 8;
        let h = 4;
        let mut semantic = vec![SemanticClass::Void as u8; w * h];
        let mut depth = vec![0i32; w * h];

        // Fill some pattern
        for x in 0..w {
            // Bottom row = floor
            semantic[3 * w + x] = SemanticClass::Floor as u8;
            depth[3 * w + x] = 100 * (x as i32 + 1);
            // Middle rows = wall
            semantic[1 * w + x] = SemanticClass::Wall as u8;
            semantic[2 * w + x] = SemanticClass::Wall as u8;
            depth[1 * w + x] = 200;
            depth[2 * w + x] = 200;
            // Top row = ceiling
            semantic[0 * w + x] = SemanticClass::Ceiling as u8;
        }
        // Put an enemy in the center
        semantic[1 * w + 4] = SemanticClass::Enemy as u8;
        depth[1 * w + 4] = 50;

        PerceptionFrame {
            tick: 0,
            width: w,
            height: h,
            semantic,
            depth,
        }
    }

    #[test]
    fn columns_with_class_finds_enemy() {
        let frame = make_frame();
        let cols = frame.columns_with_class(SemanticClass::Enemy);
        assert_eq!(cols, vec![4]);
    }

    #[test]
    fn nearest_depth_of_enemy() {
        let frame = make_frame();
        let d = frame.nearest_depth_of(SemanticClass::Enemy);
        assert_eq!(d, Some(50));
    }

    #[test]
    fn nearest_depth_of_absent() {
        let frame = make_frame();
        let d = frame.nearest_depth_of(SemanticClass::Player);
        assert_eq!(d, None);
    }

    #[test]
    fn downsample_halves() {
        let frame = make_frame();
        let small = frame.downsample(4, 2);
        assert_eq!(small.width, 4);
        assert_eq!(small.height, 2);
        assert_eq!(small.semantic.len(), 8);
        assert_eq!(small.depth.len(), 8);
    }

    #[test]
    fn count_class_wall() {
        let frame = make_frame();
        // 2 rows × 8 cols = 16, minus 1 enemy pixel = 15
        assert_eq!(frame.count_class(SemanticClass::Wall), 15);
    }
}