//! HTMRegion: compose SpatialPooler + TemporalMemory into a single step().

use crate::sp::{SpatialPooler, SpatialPoolerConfig};
use crate::tm::{TemporalMemory, TemporalMemoryConfig};
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize)]
pub struct HTMRegionCore {
    pub sp: SpatialPooler,
    pub tm: TemporalMemory,
}

impl HTMRegionCore {
    pub fn new(
        input_bits: usize,
        n_columns: usize,
        cells_per_column: usize,
        seed: u64,
    ) -> Self {
        let defaults = SpatialPoolerConfig::default();
        let sp_cfg = SpatialPoolerConfig {
            input_bits,
            n_columns,
            // Scale potential_radius to at most the input size.
            potential_radius: defaults.potential_radius.min(input_bits),
            ..defaults
        };

        let tm_cfg = TemporalMemoryConfig {
            n_columns,
            cells_per_column,
            ..TemporalMemoryConfig::default()
        };

        Self {
            sp: SpatialPooler::new(sp_cfg, seed),
            tm: TemporalMemory::new(tm_cfg, seed.wrapping_add(0x9E3779B97F4A7C15)),
        }
    }

    /// Process one timestep. Returns (active_columns_mask,
    /// active_cells_mask, predicted_cells_mask, anomaly).
    pub fn step(
        &mut self,
        input_sdr: &[bool],
        learn: bool,
    ) -> (Vec<bool>, Vec<bool>, Vec<bool>, f32) {
        let active_cols = self.sp.compute(input_sdr, learn);

        let mut active_cols_mask = vec![false; self.sp.cfg.n_columns];
        for &c in &active_cols {
            active_cols_mask[c as usize] = true;
        }

        let anomaly = self.tm.compute(&active_cols, learn);

        // active_cells and predictive_cells are stored as Vec<bool> already.
        let active_cells_mask = self.tm.active_cells.clone();
        let predicted_cells_mask = self.tm.predictive_cells.clone();

        (active_cols_mask, active_cells_mask, predicted_cells_mask, anomaly)
    }

    pub fn reset(&mut self) {
        self.tm.reset();
    }

    /// Process T timesteps in one call. Returns flat `(T*n_columns)` active-column
    /// mask (u8 0/1) and `(T,)` anomaly scores.
    ///
    /// Amortises the per-step Python round-trip for training: one GIL release,
    /// one copy-out. Used by `HTMLayer.step_many`.
    pub fn step_many(
        &mut self,
        inputs_flat: &[bool],
        input_bits: usize,
        t: usize,
        learn: bool,
    ) -> (Vec<u8>, Vec<f32>) {
        let n_cols = self.sp.cfg.n_columns;
        debug_assert_eq!(inputs_flat.len(), t * input_bits);
        let mut cols = vec![0u8; t * n_cols];
        let mut anom = vec![0f32; t];
        for ti in 0..t {
            let off = ti * input_bits;
            let input = &inputs_flat[off..off + input_bits];
            let active_cols = self.sp.compute(input, learn);
            let co = ti * n_cols;
            for &c in &active_cols {
                cols[co + c as usize] = 1;
            }
            anom[ti] = self.tm.compute(&active_cols, learn);
        }
        (cols, anom)
    }
}