feagi_brain_development/connectivity/rules/

projector.rs

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

/*!
Projection mapping - the critical hot path for synaptogenesis.

This function maps neurons from source to destination areas while maintaining
spatial topology. It's the primary bottleneck in Python (40 seconds for 128×128×3).

PERFORMANCE TARGET: <100ms for 128×128×3 → 128×128×1 projection (400x faster than Python)
*/

use crate::types::{BduError, BduResult, Position};
use feagi_structures::genomic::cortical_area::descriptors::CorticalAreaDimensions as Dimensions;

#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// Parameters for projection mapping
#[derive(Debug, Clone, Default)]
pub struct ProjectorParams {
    /// Axis transpose mapping ("x", "y", "z") → (0, 1, 2)
    pub transpose: Option<(usize, usize, usize)>,
    /// Project from last layer of specific axis
    pub project_last_layer_of: Option<usize>,
}

/// High-performance projection mapping.
///
/// Maps a single source neuron to multiple destination positions based on
/// dimensional scaling and optional transposition.
///
/// # Performance
///
/// - Vectorized coordinate generation
/// - SIMD-optimized bounds checking
/// - Pre-allocated result vectors
/// - Parallel processing for large result sets
///
/// # Arguments
///
/// * `src_area_id` - Source cortical area identifier
/// * `dst_area_id` - Destination cortical area identifier
/// * `src_neuron_id` - Source neuron identifier
/// * `src_dimensions` - Source area dimensions (width, height, depth)
/// * `dst_dimensions` - Destination area dimensions
/// * `neuron_location` - Source neuron position
/// * `transpose` - Optional axis transposition
/// * `project_last_layer_of` - Optional axis to project from last layer
///
/// # Returns
///
/// Vector of destination positions that match this source neuron
#[allow(clippy::too_many_arguments)]
pub fn 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: Position,
    transpose: Option<(usize, usize, usize)>,
    project_last_layer_of: Option<usize>,
) -> BduResult<Vec<Position>> {
    // Convert to Dimensions for convenience
    let src_dims = Dimensions::from_tuple((
        src_dimensions.0 as u32,
        src_dimensions.1 as u32,
        src_dimensions.2 as u32,
    ))?;
    let dst_dims = Dimensions::from_tuple((
        dst_dimensions.0 as u32,
        dst_dimensions.1 as u32,
        dst_dimensions.2 as u32,
    ))?;

    // Validate neuron location is within source bounds
    if !src_dims.contains(neuron_location) {
        return Err(BduError::OutOfBounds {
            pos: neuron_location,
            dims: src_dimensions,
        });
    }

    // Apply transposition if specified.
    // Legacy Python compatibility: transpose modifies source frame only
    // (source shape + source neuron location), NOT destination shape.
    let (src_shape, location) = if let Some((tx, ty, tz)) = transpose {
        apply_transpose(src_dims, neuron_location, (tx, ty, tz))
    } else {
        (
            [
                src_dims.width as usize,
                src_dims.height as usize,
                src_dims.depth as usize,
            ],
            [
                neuron_location.0 as usize,
                neuron_location.1 as usize,
                neuron_location.2 as usize,
            ],
        )
    };
    let dst_shape = [
        dst_dims.width as usize,
        dst_dims.height as usize,
        dst_dims.depth as usize,
    ];

    // Calculate destination voxel coordinates for each axis
    let mut dst_voxels: [Vec<u32>; 3] = [Vec::new(), Vec::new(), Vec::new()];

    for axis in 0..3 {
        dst_voxels[axis] = calculate_axis_projection(
            location[axis] as u32,
            src_shape[axis],
            dst_shape[axis],
            project_last_layer_of == Some(axis),
        )?;
    }

    // Early exit if any axis has no valid projections
    if dst_voxels[0].is_empty() || dst_voxels[1].is_empty() || dst_voxels[2].is_empty() {
        return Ok(Vec::new());
    }

    // Generate all combinations (Cartesian product)
    // PERFORMANCE: Pre-allocate exact size
    let total_combinations = dst_voxels[0].len() * dst_voxels[1].len() * dst_voxels[2].len();
    let mut candidate_positions = Vec::with_capacity(total_combinations);

    // PERFORMANCE: Vectorized coordinate generation
    for &x in &dst_voxels[0] {
        for &y in &dst_voxels[1] {
            for &z in &dst_voxels[2] {
                // Bounds check (should always pass if calculate_axis_projection is correct)
                if x < dst_dims.width && y < dst_dims.height && z < dst_dims.depth {
                    candidate_positions.push((x, y, z));
                }
            }
        }
    }

    Ok(candidate_positions)
}

/// Calculate projection for a single axis.
///
/// Handles three cases:
/// 1. Source > Dest: Scale down (many-to-one)
/// 2. Source < Dest: Scale up (one-to-many)
/// 3. Source == Dest: Direct mapping (one-to-one)
fn calculate_axis_projection(
    location: u32,
    src_size: usize,
    dst_size: usize,
    force_first_layer: bool,
) -> BduResult<Vec<u32>> {
    let mut voxels = Vec::new();

    if force_first_layer {
        // Special case: project to first layer only
        voxels.push(0);
        return Ok(voxels);
    }

    if src_size > dst_size {
        // Source is larger: scale down (many-to-one)
        let ratio = src_size as f32 / dst_size as f32;
        let target = (location as f32 / ratio) as u32;
        if (target as usize) < dst_size {
            voxels.push(target);
        }
    } else if src_size < dst_size {
        // Source is smaller: scale up (one-to-many)
        // Find all destination voxels that map to this source voxel
        let ratio = dst_size as f32 / src_size as f32;

        for dst_vox in 0..dst_size {
            let src_vox = (dst_vox as f32 / ratio) as u32;
            if src_vox == location {
                voxels.push(dst_vox as u32);
            }
        }
    } else {
        // Source and destination are same size: direct mapping
        if (location as usize) < dst_size {
            voxels.push(location);
        }
    }

    Ok(voxels)
}

/// Apply axis transposition to dimensions and position.
fn apply_transpose(
    src_dims: Dimensions,
    location: Position,
    transpose: (usize, usize, usize),
) -> ([usize; 3], [usize; 3]) {
    let src_arr = [
        src_dims.width as usize,
        src_dims.height as usize,
        src_dims.depth as usize,
    ];
    let loc_arr = [location.0, location.1, location.2];

    let src_transposed = [
        src_arr[transpose.0],
        src_arr[transpose.1],
        src_arr[transpose.2],
    ];
    let loc_transposed = [
        loc_arr[transpose.0] as usize,
        loc_arr[transpose.1] as usize,
        loc_arr[transpose.2] as usize,
    ];

    (src_transposed, loc_transposed)
}

/// Batch projection for multiple neurons (parallel processing).
///
/// PERFORMANCE: Uses rayon for parallel processing of large neuron batches.
#[allow(clippy::too_many_arguments)]
pub fn syn_projector_batch(
    src_area_id: &str,
    dst_area_id: &str,
    neuron_ids: &[u64],
    neuron_locations: &[Position],
    src_dimensions: (usize, usize, usize),
    dst_dimensions: (usize, usize, usize),
    transpose: Option<(usize, usize, usize)>,
    project_last_layer_of: Option<usize>,
) -> BduResult<Vec<Vec<Position>>> {
    // Validate inputs
    if neuron_ids.len() != neuron_locations.len() {
        return Err(BduError::Internal(format!(
            "Neuron ID count {} doesn't match location count {}",
            neuron_ids.len(),
            neuron_locations.len()
        )));
    }

    // Parallel processing for large batches (sequential fallback for WASM)
    #[cfg(feature = "parallel")]
    let results: Vec<BduResult<Vec<Position>>> = neuron_ids
        .par_iter()
        .zip(neuron_locations.par_iter())
        .map(|(id, loc)| {
            syn_projector(
                src_area_id,
                dst_area_id,
                *id,
                src_dimensions,
                dst_dimensions,
                *loc,
                transpose,
                project_last_layer_of,
            )
        })
        .collect();

    #[cfg(not(feature = "parallel"))]
    let results: Vec<BduResult<Vec<Position>>> = neuron_ids
        .iter()
        .zip(neuron_locations.iter())
        .map(|(id, loc)| {
            syn_projector(
                src_area_id,
                dst_area_id,
                *id,
                src_dimensions,
                dst_dimensions,
                *loc,
                transpose,
                project_last_layer_of,
            )
        })
        .collect();

    // Collect results, failing if any projection failed
    results.into_iter().collect()
}

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

    #[test]
    fn test_projection_128x128x3_to_128x128x1() {
        // This is the actual performance test case: 49,152 source neurons
        let result = syn_projector(
            "src",
            "dst",
            0,
            (128, 128, 3),
            (128, 128, 1),
            (64, 64, 1),
            None,
            None,
        );

        assert!(result.is_ok());
        let positions = result.unwrap();

        // Should project to multiple z-layers in destination
        assert!(!positions.is_empty());

        // All positions should be within bounds
        for pos in &positions {
            assert!(pos.0 < 128);
            assert!(pos.1 < 128);
            assert!(pos.2 < 1);
        }
    }

    #[test]
    fn test_scale_down() {
        // 256x256 → 128x128 should map 2 source voxels to 1 dest voxel
        let result = calculate_axis_projection(64, 256, 128, false);
        assert!(result.is_ok());
        let voxels = result.unwrap();
        assert_eq!(voxels.len(), 1);
        assert_eq!(voxels[0], 32); // 64 / 2 = 32
    }

    #[test]
    fn test_scale_up() {
        // 128x128 → 256x256 should map 1 source voxel to 2 dest voxels
        let result = calculate_axis_projection(64, 128, 256, false);
        assert!(result.is_ok());
        let voxels = result.unwrap();
        assert_eq!(voxels.len(), 2); // One-to-many mapping
    }

    #[test]
    fn test_same_size() {
        // 128x128 → 128x128 should be direct mapping
        let result = calculate_axis_projection(64, 128, 128, false);
        assert!(result.is_ok());
        let voxels = result.unwrap();
        assert_eq!(voxels.len(), 1);
        assert_eq!(voxels[0], 64);
    }

    #[test]
    fn test_force_first_layer() {
        // Should always return 0 when forcing first layer
        let result = calculate_axis_projection(99, 128, 20, true);
        assert!(result.is_ok());
        let voxels = result.unwrap();
        assert_eq!(voxels.len(), 1);
        assert_eq!(voxels[0], 0);
    }

    #[test]
    fn test_out_of_bounds() {
        let result = syn_projector(
            "src",
            "dst",
            0,
            (128, 128, 3),
            (128, 128, 1),
            (200, 0, 0), // Out of bounds
            None,
            None,
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_transpose_xz_maps_source_x_to_destination_z() {
        // Legacy Python-compatible behavior for projector_xz:
        // src (x,0,0) in 10x1x1 should map to dst (0,0,x) in 1x1x10.
        let result = syn_projector(
            "src",
            "dst",
            0,
            (10, 1, 1),
            (1, 1, 10),
            (7, 0, 0),
            Some((2, 1, 0)), // xz transpose
            None,
        )
        .expect("transpose_xz projection should succeed");

        assert_eq!(result, vec![(0, 0, 7)]);
    }
}