native_neural_network 0.3.1

use crate::sphere5d::{NeuronPoint, NeuronPointF64};

pub fn layer_bounds(points: &[NeuronPoint], layer: usize) -> Option<(usize, usize)> {
    let mut first = None;
    let mut last = None;
    for (i, p) in points.iter().enumerate() {
        if p.layer == layer {
            if first.is_none() {
                first = Some(i);
            }
            last = Some(i);
        }
    }
    Some((first?, last? + 1))
}

pub fn neuron_exists(points: &[NeuronPoint], layer: usize, neuron: usize) -> bool {
    for p in points {
        if p.layer == layer && p.neuron == neuron {
            return true;
        }
    }
    false
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KernelizeError {
    InvalidDistance,
    InvalidKernelSize,
    CapacityTooSmall,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct NeuronKernel {
    pub layer: usize,
    pub start: usize,
    pub len: usize,
    pub centroid: [f32; 5],
    pub mean_bias: f32,
    pub mean_activation: f32,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct NeuronKernelF64 {
    pub layer: usize,
    pub start: usize,
    pub len: usize,
    pub centroid: [f64; 5],
    pub mean_bias: f64,
    pub mean_activation: f64,
}

impl Default for NeuronKernelF64 {
    fn default() -> Self {
        Self {
            layer: 0,
            start: 0,
            len: 0,
            centroid: [0.0; 5],
            mean_bias: 0.0,
            mean_activation: 0.0,
        }
    }
}

impl Default for NeuronKernel {
    fn default() -> Self {
        Self {
            layer: 0,
            start: 0,
            len: 0,
            centroid: [0.0; 5],
            mean_bias: 0.0,
            mean_activation: 0.0,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct KernelLayoutStats {
    pub points_used: usize,
    pub kernels_used: usize,
}

pub fn aggregate_constellation_to_kernels(
    points: &[NeuronPoint],
    max_distance: f32,
    max_kernel_size: usize,
    kernels_out: &mut [NeuronKernel],
    member_layout_out: &mut [usize],
    assignment_scratch: &mut [usize],
) -> Result<KernelLayoutStats, KernelizeError> {
    if !max_distance.is_finite() || max_distance <= 0.0 {
        return Err(KernelizeError::InvalidDistance);
    }
    if max_kernel_size == 0 {
        return Err(KernelizeError::InvalidKernelSize);
    }
    if assignment_scratch.len() < points.len() || member_layout_out.len() < points.len() {
        return Err(KernelizeError::CapacityTooSmall);
    }

    for slot in assignment_scratch.iter_mut().take(points.len()) {
        *slot = usize::MAX;
    }

    if points.is_empty() {
        return Ok(KernelLayoutStats {
            points_used: 0,
            kernels_used: 0,
        });
    }

    let mut kernels_used = 0usize;
    let mut points_used = 0usize;
    let max_d2 = max_distance * max_distance;

    let mut anchor_idx = 0usize;
    while anchor_idx < points.len() {
        if assignment_scratch[anchor_idx] != usize::MAX {
            anchor_idx += 1;
            continue;
        }

        if kernels_used >= kernels_out.len() {
            return Err(KernelizeError::CapacityTooSmall);
        }

        let anchor = points[anchor_idx];
        let start = points_used;
        let mut len = 0usize;
        let mut centroid = [0.0f32; 5];
        let mut bias_sum = 0.0f32;
        let mut activation_sum = 0.0f32;

        let mut candidate_idx = anchor_idx;
        while candidate_idx < points.len() {
            if len >= max_kernel_size {
                break;
            }
            if assignment_scratch[candidate_idx] != usize::MAX {
                candidate_idx += 1;
                continue;
            }

            let candidate = points[candidate_idx];
            if candidate.layer != anchor.layer {
                candidate_idx += 1;
                continue;
            }

            let mut d2 = 0.0f32;
            let mut axis = 0usize;
            while axis < 5 {
                let d = candidate.position[axis] - anchor.position[axis];
                d2 += d * d;
                axis += 1;
            }
            if d2 > max_d2 {
                candidate_idx += 1;
                continue;
            }

            member_layout_out[points_used] = candidate_idx;
            assignment_scratch[candidate_idx] = kernels_used;
            points_used += 1;
            len += 1;

            let mut k = 0usize;
            while k < 5 {
                centroid[k] += candidate.position[k];
                k += 1;
            }
            bias_sum += candidate.bias;
            activation_sum += candidate.activation;

            candidate_idx += 1;
        }

        if len == 0 {
            member_layout_out[points_used] = anchor_idx;
            assignment_scratch[anchor_idx] = kernels_used;
            points_used += 1;
            len = 1;
            centroid = anchor.position;
            bias_sum = anchor.bias;
            activation_sum = anchor.activation;
        } else {
            let inv_len = 1.0f32 / (len as f32);
            let mut k = 0usize;
            while k < 5 {
                centroid[k] *= inv_len;
                k += 1;
            }
            bias_sum *= inv_len;
            activation_sum *= inv_len;
        }

        kernels_out[kernels_used] = NeuronKernel {
            layer: anchor.layer,
            start,
            len,
            centroid,
            mean_bias: bias_sum,
            mean_activation: activation_sum,
        };
        kernels_used += 1;
        anchor_idx += 1;
    }

    Ok(KernelLayoutStats {
        points_used,
        kernels_used,
    })
}

pub fn aggregate_constellation_to_kernels_f64(
    points: &[NeuronPointF64],
    max_distance: f64,
    max_kernel_size: usize,
    kernels_out: &mut [NeuronKernelF64],
    member_layout_out: &mut [usize],
    assignment_scratch: &mut [usize],
) -> Result<KernelLayoutStats, KernelizeError> {
    if !max_distance.is_finite() || max_distance <= 0.0 {
        return Err(KernelizeError::InvalidDistance);
    }
    if max_kernel_size == 0 {
        return Err(KernelizeError::InvalidKernelSize);
    }
    if assignment_scratch.len() < points.len() || member_layout_out.len() < points.len() {
        return Err(KernelizeError::CapacityTooSmall);
    }

    for slot in assignment_scratch.iter_mut().take(points.len()) {
        *slot = usize::MAX;
    }

    if points.is_empty() {
        return Ok(KernelLayoutStats {
            points_used: 0,
            kernels_used: 0,
        });
    }

    let mut kernels_used = 0usize;
    let mut points_used = 0usize;
    let max_d2 = max_distance * max_distance;

    let mut anchor_idx = 0usize;
    while anchor_idx < points.len() {
        if assignment_scratch[anchor_idx] != usize::MAX {
            anchor_idx += 1;
            continue;
        }

        if kernels_used >= kernels_out.len() {
            return Err(KernelizeError::CapacityTooSmall);
        }

        let anchor = points[anchor_idx];
        let start = points_used;
        let mut len = 0usize;
        let mut centroid = [0.0f64; 5];
        let mut bias_sum = 0.0f64;
        let mut activation_sum = 0.0f64;

        let mut candidate_idx = anchor_idx;
        while candidate_idx < points.len() {
            if len >= max_kernel_size {
                break;
            }
            if assignment_scratch[candidate_idx] != usize::MAX {
                candidate_idx += 1;
                continue;
            }

            let candidate = points[candidate_idx];
            if candidate.layer != anchor.layer {
                candidate_idx += 1;
                continue;
            }

            let mut d2 = 0.0f64;
            let mut axis = 0usize;
            while axis < 5 {
                let d = candidate.position[axis] - anchor.position[axis];
                d2 += d * d;
                axis += 1;
            }
            if d2 > max_d2 {
                candidate_idx += 1;
                continue;
            }

            member_layout_out[points_used] = candidate_idx;
            assignment_scratch[candidate_idx] = kernels_used;
            points_used += 1;
            len += 1;

            let mut k = 0usize;
            while k < 5 {
                centroid[k] += candidate.position[k];
                k += 1;
            }
            bias_sum += candidate.bias;
            activation_sum += candidate.activation;

            candidate_idx += 1;
        }

        if len == 0 {
            member_layout_out[points_used] = anchor_idx;
            assignment_scratch[anchor_idx] = kernels_used;
            points_used += 1;
            len = 1;
            centroid = anchor.position;
            bias_sum = anchor.bias;
            activation_sum = anchor.activation;
        } else {
            let inv_len = 1.0f64 / (len as f64);
            let mut k = 0usize;
            while k < 5 {
                centroid[k] *= inv_len;
                k += 1;
            }
            bias_sum *= inv_len;
            activation_sum *= inv_len;
        }

        kernels_out[kernels_used] = NeuronKernelF64 {
            layer: anchor.layer,
            start,
            len,
            centroid,
            mean_bias: bias_sum,
            mean_activation: activation_sum,
        };
        kernels_used += 1;
        anchor_idx += 1;
    }

    Ok(KernelLayoutStats {
        points_used,
        kernels_used,
    })
}