feagi-brain-development 0.0.22

// Copyright 2025 Neuraville Inc.
// SPDX-License-Identifier: Apache-2.0

/*!
FFI bindings for Python integration via PyO3.

Exposes Rust functions to Python for seamless interop.
*/

use pyo3::prelude::*;

/// Python-exposed version of syn_projector
///
/// # Arguments (from Python)
///
/// * `src_area_id` - str
/// * `dst_area_id` - str
/// * `src_neuron_id` - int
/// * `src_dimensions` - tuple[int, int, int]
/// * `dst_dimensions` - tuple[int, int, int]
/// * `neuron_location` - tuple[int, int, int]
/// * `transpose` - Optional[tuple[int, int, int]]
/// * `project_last_layer_of` - Optional[int]
///
/// # Returns
///
/// List[tuple[int, int, int]] - List of destination positions
#[pyfunction]
#[pyo3(signature = (src_area_id, dst_area_id, src_neuron_id, src_dimensions, dst_dimensions, neuron_location, transpose=None, project_last_layer_of=None))]
fn py_syn_projector(
    src_area_id: &str,
    dst_area_id: &str,
    src_neuron_id: u64,
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
    neuron_location: (i32, i32, i32),
    transpose: Option<(usize, usize, usize)>,
    project_last_layer_of: Option<usize>,
) -> PyResult<Vec<(i32, i32, i32)>> {
    let loc_u32 = (neuron_location.0 as u32, neuron_location.1 as u32, neuron_location.2 as u32);
    let result = crate::connectivity::rules::syn_projector(
        src_area_id,
        dst_area_id,
        src_neuron_id,
        src_dimensions,
        dst_dimensions,
        loc_u32,
        transpose,
        project_last_layer_of,
    );

    match result {
        Ok(positions) => Ok(positions.iter().map(|(x, y, z)| (*x as i32, *y as i32, *z as i32)).collect()),
        Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
            "Rust syn_projector error: {}",
            e
        ))),
    }
}

/// Python-exposed batch projector for parallel processing
///
/// # Arguments
///
/// * `src_area_id` - str
/// * `dst_area_id` - str
/// * `neuron_ids` - List[int]
/// * `neuron_locations` - List[tuple[int, int, int]]
/// * `src_dimensions` - tuple[int, int, int]
/// * `dst_dimensions` - tuple[int, int, int]
/// * `transpose` - Optional[tuple[int, int, int]]
/// * `project_last_layer_of` - Optional[int]
///
/// # Returns
///
/// List[List[tuple[int, int, int]]] - List of position lists (one per neuron)
#[pyfunction]
#[pyo3(signature = (src_area_id, dst_area_id, neuron_ids, neuron_locations, src_dimensions, dst_dimensions, transpose=None, project_last_layer_of=None))]
fn py_syn_projector_batch(
    src_area_id: &str,
    dst_area_id: &str,
    neuron_ids: Vec<u64>,
    neuron_locations: Vec<(i32, i32, i32)>,
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
    transpose: Option<(usize, usize, usize)>,
    project_last_layer_of: Option<usize>,
) -> PyResult<Vec<Vec<(i32, i32, i32)>>> {
    let locs_u32: Vec<(u32, u32, u32)> = neuron_locations.iter()
        .map(|(x, y, z)| (*x as u32, *y as u32, *z as u32))
        .collect();
    let result = crate::connectivity::rules::syn_projector_batch(
        src_area_id,
        dst_area_id,
        &neuron_ids,
        &locs_u32,
        src_dimensions,
        dst_dimensions,
        transpose,
        project_last_layer_of,
    );

    match result {
        Ok(position_lists) => Ok(position_lists.iter()
            .map(|positions| positions.iter()
                .map(|(x, y, z)| (*x as i32, *y as i32, *z as i32))
                .collect())
            .collect()),
        Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
            "Rust syn_projector_batch error: {}",
            e
        ))),
    }
}

/// Block connection - maps blocks with scaling factor
#[pyfunction]
fn py_syn_block_connection(
    src_area_id: &str,
    dst_area_id: &str,
    neuron_location: (i32, i32, i32),
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
    scaling_factor: i32,
) -> PyResult<(i32, i32, i32)> {
    // Convert from Python's i32 coordinates to Rust's u32 coordinates
    let loc_u32 = (neuron_location.0 as u32, neuron_location.1 as u32, neuron_location.2 as u32);
    let scale_u32 = scaling_factor as u32;

    match crate::connectivity::rules::syn_block_connection(
        src_area_id, dst_area_id, loc_u32, src_dimensions, dst_dimensions, scale_u32
    ) {
        Ok((x, y, z)) => Ok((x as i32, y as i32, z as i32)),
        Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{}", e))),
    }
}

/// Expander - scales coordinates from source to destination
#[pyfunction]
fn py_syn_expander(
    src_area_id: &str,
    dst_area_id: &str,
    neuron_location: (i32, i32, i32),
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
) -> PyResult<(i32, i32, i32)> {
    let loc_u32 = (neuron_location.0 as u32, neuron_location.1 as u32, neuron_location.2 as u32);
    match crate::connectivity::rules::syn_expander(
        src_area_id, dst_area_id, loc_u32, src_dimensions, dst_dimensions
    ) {
        Ok((x, y, z)) => Ok((x as i32, y as i32, z as i32)),
        Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{}", e))),
    }
}

/// Expander batch - parallel processing
#[pyfunction]
fn py_syn_expander_batch(
    src_area_id: &str,
    dst_area_id: &str,
    neuron_locations: Vec<(i32, i32, i32)>,
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
) -> PyResult<Vec<(i32, i32, i32)>> {
    let locs_u32: Vec<(u32, u32, u32)> = neuron_locations.iter()
        .map(|(x, y, z)| (*x as u32, *y as u32, *z as u32))
        .collect();
    match crate::connectivity::rules::syn_expander_batch(
        src_area_id, dst_area_id, &locs_u32, src_dimensions, dst_dimensions
    ) {
        Ok(positions) => Ok(positions.iter().map(|(x, y, z)| (*x as i32, *y as i32, *z as i32)).collect()),
        Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{}", e))),
    }
}

/// Reducer - binary encoding to multiple positions
#[pyfunction]
fn py_syn_reducer_x(
    src_area_id: &str,
    dst_area_id: &str,
    neuron_location: (i32, i32, i32),
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
    dst_y_index: i32,
    dst_z_index: i32,
) -> PyResult<Vec<(i32, i32, i32)>> {
    let loc_u32 = (neuron_location.0 as u32, neuron_location.1 as u32, neuron_location.2 as u32);
    match crate::connectivity::rules::syn_reducer_x(
        src_area_id, dst_area_id, loc_u32, src_dimensions, dst_dimensions, dst_y_index as u32, dst_z_index as u32
    ) {
        Ok(positions) => Ok(positions.iter().map(|(x, y, z)| (*x as i32, *y as i32, *z as i32)).collect()),
        Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{}", e))),
    }
}

/// Python wrapper for MortonSpatialHash
#[pyclass]
struct PyMortonSpatialHash {
    inner: std::sync::Arc<crate::spatial::MortonSpatialHash>,
}

#[pymethods]
impl PyMortonSpatialHash {
    #[new]
    fn new() -> Self {
        Self {
            inner: std::sync::Arc::new(crate::spatial::MortonSpatialHash::new()),
        }
    }

    fn add_neuron(&self, cortical_area: String, x: u32, y: u32, z: u32, neuron_id: u64) -> bool {
        self.inner.add_neuron(cortical_area, x, y, z, neuron_id)
    }

    fn get_neuron_at_coordinate(&self, cortical_area: &str, x: u32, y: u32, z: u32) -> Option<u64> {
        self.inner.get_neuron_at_coordinate(cortical_area, x, y, z)
    }

    fn get_neurons_at_coordinate(&self, cortical_area: &str, x: u32, y: u32, z: u32) -> Vec<u64> {
        self.inner.get_neurons_at_coordinate(cortical_area, x, y, z)
    }

    fn get_neurons_in_region(
        &self,
        cortical_area: &str,
        x1: u32, y1: u32, z1: u32,
        x2: u32, y2: u32, z2: u32,
    ) -> Vec<u64> {
        self.inner.get_neurons_in_region(cortical_area, x1, y1, z1, x2, y2, z2)
    }

    fn get_neuron_position(&self, neuron_id: u64) -> Option<(String, u32, u32, u32)> {
        self.inner.get_neuron_position(neuron_id)
    }

    fn remove_neuron(&self, neuron_id: u64) -> bool {
        self.inner.remove_neuron(neuron_id)
    }

    fn clear(&self) {
        self.inner.clear();
    }

    fn get_stats(&self) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            let stats = self.inner.get_stats();
            let dict = pyo3::types::PyDict::new_bound(py);
            dict.set_item("total_areas", stats.total_areas)?;
            dict.set_item("total_neurons", stats.total_neurons)?;
            dict.set_item("total_occupied_positions", stats.total_occupied_positions)?;
            Ok(dict.to_object(py))
        })
    }
}

/// Morton encode 3D coordinates
#[pyfunction]
fn py_morton_encode_3d(x: u32, y: u32, z: u32) -> PyResult<u64> {
    use crate::spatial::morton_encode_3d;
    Ok(morton_encode_3d(x, y, z))
}

/// Morton decode to 3D coordinates
#[pyfunction]
fn py_morton_decode_3d(morton_code: u64) -> PyResult<(u32, u32, u32)> {
    use crate::spatial::morton_decode_3d;
    Ok(morton_decode_3d(morton_code))
}

/// Vector offset - apply single vector to position
#[pyfunction]
fn py_apply_vector_offset(
    src_position: (i32, i32, i32),
    vector: (i32, i32, i32),
    morphology_scalar: f32,
    dst_dimensions: (usize, usize, usize),
) -> PyResult<Option<(i32, i32, i32)>> {
    let src_u32 = (src_position.0 as u32, src_position.1 as u32, src_position.2 as u32);
    Ok(
        crate::connectivity::rules::apply_vector_offset(src_u32, vector, morphology_scalar, dst_dimensions)
            .map(|(x, y, z)| (x as i32, y as i32, z as i32)),
    )
}

/// Vector batch - apply vector to multiple positions
#[pyfunction]
fn py_match_vectors_batch(
    src_positions: Vec<(i32, i32, i32)>,
    vector: (i32, i32, i32),
    morphology_scalar: f32,
    dst_dimensions: (usize, usize, usize),
) -> PyResult<Vec<(i32, i32, i32)>> {
    let srcs_u32: Vec<(u32, u32, u32)> = src_positions.iter()
        .map(|(x, y, z)| (*x as u32, *y as u32, *z as u32))
        .collect();
    match crate::connectivity::rules::match_vectors_batch(
        &srcs_u32, vector, morphology_scalar, dst_dimensions
    ) {
        Ok(positions) => Ok(positions.iter().map(|(x, y, z)| (*x as i32, *y as i32, *z as i32)).collect()),
        Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!("{}", e))),
    }
}

/// Pattern matching - find destinations from pattern rules
/// patterns: list of (src_pattern, dst_pattern) tuples
/// Each pattern is (x, y, z) where elements can be:
///   -1 = wildcard "*"
///   -2 = skip "?"
///   -3 = exclude "!"
///   -10 = direction positive "?+" (all coords > src)
///   -11 = direction negative "?-" (all coords < src)
///   -12 = direction positive inclusive "?+=" (all coords >= src)
///   -13 = direction negative inclusive "?-=" (all coords <= src)
///   >= 0 = exact value
/// Note: Offset and Range patterns require string-based genome encoding.
#[pyfunction]
fn py_match_patterns(
    src_coordinate: (i32, i32, i32),
    patterns: Vec<((i32, i32, i32), (i32, i32, i32))>,
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
) -> PyResult<Vec<(i32, i32, i32)>> {
    use crate::connectivity::rules::patterns::{PatternElement, match_patterns_batch};

    let src_u32 = (src_coordinate.0 as u32, src_coordinate.1 as u32, src_coordinate.2 as u32);

    // Convert integer patterns to PatternElement
    let parsed_patterns: Vec<_> = patterns.iter().map(|(src, dst)| {
        let src_pattern = (
            PatternElement::from_int(src.0),
            PatternElement::from_int(src.1),
            PatternElement::from_int(src.2),
        );
        let dst_pattern = (
            PatternElement::from_int(dst.0),
            PatternElement::from_int(dst.1),
            PatternElement::from_int(dst.2),
        );
        (src_pattern, dst_pattern)
    }).collect();

    let results = match_patterns_batch(
        src_u32,
        &parsed_patterns,
        src_dimensions,
        dst_dimensions,
    );

    Ok(results.iter().map(|(x, y, z)| (*x as i32, *y as i32, *z as i32)).collect())
}

/// Find source coordinates that match a pattern
#[pyfunction]
fn py_find_source_coordinates(
    src_pattern: (i32, i32, i32),
    src_dimensions: (usize, usize, usize),
) -> PyResult<Vec<(i32, i32, i32)>> {
    use crate::connectivity::rules::patterns::{PatternElement, find_source_coordinates};

    let pattern = (
        PatternElement::from_int(src_pattern.0),
        PatternElement::from_int(src_pattern.1),
        PatternElement::from_int(src_pattern.2),
    );

    let results = find_source_coordinates(&pattern, src_dimensions);
    Ok(results.iter().map(|(x, y, z)| (*x as i32, *y as i32, *z as i32)).collect())
}

/// Randomizer - select random position in destination area
#[pyfunction]
fn py_syn_randomizer(dst_dimensions: (usize, usize, usize)) -> PyResult<(i32, i32, i32)> {
    let (x, y, z) = crate::connectivity::rules::syn_randomizer(dst_dimensions);
    Ok((x as i32, y as i32, z as i32))
}

/// Lateral pairs X - connect neurons in pairs along X axis
#[pyfunction]
fn py_syn_lateral_pairs_x(
    neuron_location: (i32, i32, i32),
    src_dimensions: (usize, usize, usize),
) -> PyResult<Option<(i32, i32, i32)>> {
    let loc_u32 = (neuron_location.0 as u32, neuron_location.1 as u32, neuron_location.2 as u32);
    Ok(crate::connectivity::rules::syn_lateral_pairs_x(loc_u32, src_dimensions)
        .map(|(x, y, z)| (x as i32, y as i32, z as i32)))
}

/// Last to first - connect last neuron to first (feedback connection)
#[pyfunction]
fn py_syn_last_to_first(
    neuron_location: (i32, i32, i32),
    src_dimensions: (usize, usize, usize),
) -> PyResult<Option<(i32, i32, i32)>> {
    let loc_u32 = (neuron_location.0 as u32, neuron_location.1 as u32, neuron_location.2 as u32);
    Ok(crate::connectivity::rules::syn_last_to_first(loc_u32, src_dimensions)
        .map(|(x, y, z)| (x as i32, y as i32, z as i32)))
}

/// Python module initialization (PyO3 0.22 API with Bound)
#[pymodule]
fn feagi_brain_development(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Projector functions
    m.add_function(wrap_pyfunction!(py_syn_projector, m)?)?;
    m.add_function(wrap_pyfunction!(py_syn_projector_batch, m)?)?;

    // Phase 2 morphologies
    m.add_function(wrap_pyfunction!(py_syn_block_connection, m)?)?;
    m.add_function(wrap_pyfunction!(py_syn_expander, m)?)?;
    m.add_function(wrap_pyfunction!(py_syn_expander_batch, m)?)?;
    m.add_function(wrap_pyfunction!(py_syn_reducer_x, m)?)?;

    // Phase 3C: Vector morphologies
    m.add_function(wrap_pyfunction!(py_apply_vector_offset, m)?)?;
    m.add_function(wrap_pyfunction!(py_match_vectors_batch, m)?)?;

    // Phase 3D: Pattern morphologies
    m.add_function(wrap_pyfunction!(py_match_patterns, m)?)?;
    m.add_function(wrap_pyfunction!(py_find_source_coordinates, m)?)?;

    // Phase 3E: Trivial morphologies
    m.add_function(wrap_pyfunction!(py_syn_randomizer, m)?)?;
    m.add_function(wrap_pyfunction!(py_syn_lateral_pairs_x, m)?)?;
    m.add_function(wrap_pyfunction!(py_syn_last_to_first, m)?)?;

    // Phase 3B: Morton spatial hash
    m.add_class::<PyMortonSpatialHash>()?;
    m.add_function(wrap_pyfunction!(py_morton_encode_3d, m)?)?;
    m.add_function(wrap_pyfunction!(py_morton_decode_3d, m)?)?;

    // Version info
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    m.add("__doc__", "FEAGI BDU - High-performance Brain Development Utilities")?;

    Ok(())
}

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

    #[test]
    fn test_ffi_compatibility() {
        // Smoke test to ensure FFI functions compile
        let result = py_syn_projector(
            "src",
            "dst",
            42,
            (128, 128, 3),
            (128, 128, 1),
            (64, 64, 1),
            None,
            None,
        );
        assert!(result.is_ok());
    }
}