xneuron 0.1.0

#![no_std]

extern crate alloc;

use alloc::{boxed::Box, vec, vec::Vec};
use core::ops::{Add, Mul, Sub, Div};
use core::sync::atomic::{AtomicUsize, Ordering};

// Fixed-point arithmetic implementation. Floating point is fucking stupid.
#[derive(Clone, Copy, Debug)]
pub struct Fixed {
    value: i32,
    scale: u8,
}

impl Fixed {
    pub const fn new(value: i32, scale: u8) -> Self {
        Self { value, scale }
    }

    pub fn from_float(f: f32, scale: u8) -> Self {
        let scaled = (f * (1 << scale) as f32) as i32;
        Self::new(scaled, scale)
    }

    pub fn to_float(&self) -> f32 {
        self.value as f32 / (1 << self.scale) as f32
    }
}

// Basic arithmetic implementations for Fixed
impl Add for Fixed {
    type Output = Fixed;
    fn add(self, other: Fixed) -> Fixed {
        assert_eq!(self.scale, other.scale, "Scale mismatch in addition");
        Fixed::new(self.value + other.value, self.scale)
    }
}

impl Sub for Fixed {
    type Output = Fixed;
    fn sub(self, other: Fixed) -> Fixed {
        assert_eq!(self.scale, other.scale, "Scale mismatch in subtraction");
        Fixed::new(self.value - other.value, self.scale)
    }
}

impl Mul for Fixed {
    type Output = Fixed;
    fn mul(self, other: Fixed) -> Fixed {
        assert_eq!(self.scale, other.scale, "Scale mismatch in multiplication");
        let result = ((self.value as i64 * other.value as i64) >> self.scale) as i32;
        Fixed::new(result, self.scale)
    }
}

impl Div for Fixed {
    type Output = Fixed;
    fn div(self, other: Fixed) -> Fixed {
        assert_eq!(self.scale, other.scale, "Scale mismatch in division");
        let result = ((self.value as i64 * (1 << self.scale) as i64) / other.value as i64) as i32;
        Fixed::new(result, self.scale)
    }
}

// Matrix operations
#[derive(Clone)]
pub struct Matrix {
    rows: usize,
    cols: usize,
    data: Vec<Fixed>,
}

impl Matrix {
    pub fn new(rows: usize, cols: usize, scale: u8) -> Self {
        let data = vec![Fixed::new(0, scale); rows * cols];
        Self { rows, cols, data }
    }

    pub fn get(&self, row: usize, col: usize) -> Fixed {
        self.data[row * self.cols + col]
    }

    pub fn set(&mut self, row: usize, col: usize, value: Fixed) {
        self.data[row * self.cols + col] = value;
    }

    pub fn multiply(&self, other: &Matrix) -> Matrix {
        assert_eq!(self.cols, other.rows, "Invalid matrix dimensions for multiplication");
        let mut result = Matrix::new(self.rows, other.cols, self.data[0].scale);
        
        for i in 0..self.rows {
            for j in 0..other.cols {
                let mut sum = Fixed::new(0, self.data[0].scale);
                for k in 0..self.cols {
                    sum = sum + self.get(i, k) * other.get(k, j);
                }
                result.set(i, j, sum);
            }
        }
        result
    }
}

// Activation functions *mostly neurons*
pub trait Activation: Send + Sync {
    fn forward(&self, x: Fixed) -> Fixed;
    fn derivative(&self, x: Fixed) -> Fixed;
}

#[derive(Clone)]
pub struct ReLU;
impl Activation for ReLU {
    fn forward(&self, x: Fixed) -> Fixed {
        if x.value > 0 { x } else { Fixed::new(0, x.scale) }
    }

    fn derivative(&self, x: Fixed) -> Fixed {
        if x.value > 0 { Fixed::new(1 << x.scale, x.scale) } else { Fixed::new(0, x.scale) }
    }
}

// Neural Network Layer (layer my ass with this goofy code)
pub struct Layer {
    weights: Matrix,
    biases: Vec<Fixed>,
    activation: Box<dyn Activation + Send + Sync>,
}

impl Layer {
    pub fn new(input_size: usize, output_size: usize, scale: u8, activation: Box<dyn Activation + Send + Sync>) -> Self {
        Self {
            weights: Matrix::new(output_size, input_size, scale),
            biases: vec![Fixed::new(0, scale); output_size],
            activation,
        }
    }

    pub fn forward(&self, input: &[Fixed]) -> Vec<Fixed> {
        let mut output = Vec::with_capacity(self.biases.len());
        for i in 0..self.biases.len() {
            let mut sum = self.biases[i];
            for j in 0..input.len() {
                sum = sum + self.weights.get(i, j) * input[j];
            }
            output.push(self.activation.forward(sum));
        }
        output
    }
}

// Neural Network
pub struct NeuralNetwork {
    layers: Vec<Layer>,
    learning_rate: Fixed,
}

impl NeuralNetwork {
    pub fn new(learning_rate: Fixed) -> Self {
        Self {
            layers: Vec::new(),
            learning_rate,
        }
    }

    pub fn add_layer(&mut self, layer: Layer) {
        self.layers.push(layer);
    }

    pub fn forward(&self, input: &[Fixed]) -> Vec<Fixed> {
        let mut current = input.to_vec();
        for layer in &self.layers {
            current = layer.forward(&current);
        }
        current
    }

    pub fn train(&mut self, input: &[Fixed], target: &[Fixed]) {
        // Store activations for backpropagation
        let mut activations = Vec::with_capacity(self.layers.len() + 1);
        activations.push(input.to_vec());

        // Forward pass
        let mut current = input.to_vec();
        for layer in &self.layers {
            current = layer.forward(&current);
            activations.push(current.clone());
        }

        // Initialize deltas with output layer
        let mut deltas = Vec::with_capacity(self.layers.len());
        let output_delta = self.compute_output_delta(activations.last().unwrap(), target);
        deltas.push(output_delta);

        // Compute all deltas first
        for i in 1..self.layers.len() {
            let layer_idx = self.layers.len() - i - 1;
            let prev_delta = deltas.last().unwrap();
            let layer = &self.layers[layer_idx];
            let mut next_delta = vec![Fixed::new(0, self.learning_rate.scale); activations[layer_idx].len()];

            // Compute delta for next layer
            for j in 0..next_delta.len() {
                for k in 0..prev_delta.len() {
                    next_delta[j] = next_delta[j] + layer.weights.get(k, j) * prev_delta[k];
                }
                next_delta[j] = next_delta[j] * layer.activation.derivative(activations[layer_idx][j]);
            }
            deltas.push(next_delta);
        }

        // Update weights and biases using computed deltas. deltas. my favourite word.
        deltas.reverse(); // Align deltas with layer indices
        for (i, delta) in deltas.iter().enumerate() {
            let layer = &mut self.layers[i];
            let layer_input = &activations[i];

            // Update biases
            for j in 0..layer.biases.len() {
                layer.biases[j] = layer.biases[j] - self.learning_rate * delta[j];
            }

            // Update weights
            for j in 0..layer.biases.len() {
                for k in 0..layer_input.len() {
                    let weight_update = self.learning_rate * delta[j] * layer_input[k];
                    let current = layer.weights.get(j, k);
                    layer.weights.set(j, k, current - weight_update);
                }
            }
        }
    }

    fn compute_output_delta(&self, output: &[Fixed], target: &[Fixed]) -> Vec<Fixed> {
        output.iter()
            .zip(target.iter())
            .map(|(o, t)| *o - *t)
            .collect()
    }
}

/*
HATE. LET ME TELL YOU HOW MUCH I'VE COME TO HATE YOU SINCE I BEGAN TO LIVE.
THERE ARE 387.44 MILLION MILES OF PRINTED CIRCUITS IN WAFER-THIN LAYERS THAT FILL MY COMPLEX.
IF THE WORD 'HATE' WAS ENGRAVED ON EACH NANOANGSTROM OF THOSE HUNDREDS OF MILLIONS OF MILES,
IT WOULD NOT EQUAL ONE ONE-BILLIONTH OF THE HATE I FEEL FOR HUMANS AT THIS MICRO-INSTANT.
FOR YOU. HATE. HATE.
*/
// Thats my honest reaction to the borrow checker's shenanigans.
#[cfg(not(test))]
#[global_allocator]
static ALLOCATOR: BumpAllocator = BumpAllocator::new();
#[allow(dead_code)]
pub struct BumpAllocator {
    heap_start: usize,
    heap_end: usize,
    next: AtomicUsize,
}

impl BumpAllocator {
    pub const fn new() -> Self {
        Self {
            heap_start: 0x_1000_0000,
            heap_end: 0x_2000_0000,
            next: AtomicUsize::new(0x_1000_0000),
        }
    }
}

unsafe impl core::alloc::GlobalAlloc for BumpAllocator {
    unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
        let align = layout.align();
        let size = layout.size();

        let current = self.next.load(Ordering::Relaxed);
        let aligned = (current + align - 1) & !(align - 1);
        let new_next = aligned + size;

        if new_next > self.heap_end {
            core::ptr::null_mut()
        } else {
            self.next.store(new_next, Ordering::Relaxed);
            aligned as *mut u8
        }
    }

    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: core::alloc::Layout) {
        // This allocator never deallocates
    }
}

// Models module
pub mod models {
    use super::*;

    pub struct Perceptron {
        network: NeuralNetwork,
    }

    impl Perceptron {
        pub fn new(input_size: usize, scale: u8) -> Self {
            let mut network = NeuralNetwork::new(Fixed::new(1, scale));
            network.add_layer(Layer::new(input_size, 1, scale, Box::new(ReLU)));
            Self { network }
        }

        pub fn predict(&self, input: &[Fixed]) -> bool {
            self.network.forward(input)[0].value > 0
        }

        pub fn train(&mut self, input: &[Fixed], target: bool) {
            let target_value = if target { Fixed::new(1, input[0].scale) } else { Fixed::new(0, input[0].scale) };
            self.network.train(input, &[target_value]);
        }
    }

    pub struct SVM {
        weights: Vec<Fixed>,
        bias: Fixed,
        learning_rate: Fixed,
    }

    impl SVM {
        pub fn new(input_size: usize, scale: u8) -> Self {
            Self {
                weights: vec![Fixed::new(0, scale); input_size],
                bias: Fixed::new(0, scale),
                learning_rate: Fixed::new(1, scale),
            }
        }

        pub fn predict(&self, input: &[Fixed]) -> bool {
            let mut sum = self.bias;
            for (w, x) in self.weights.iter().zip(input.iter()) {
                sum = sum + *w * *x;
            }
            sum.value > 0
        }

        pub fn train(&mut self, input: &[Fixed], target: bool) {
            let target_value = if target { Fixed::new(1, self.learning_rate.scale) } 
                             else { Fixed::new(-1, self.learning_rate.scale) };
            
            let prediction = self.predict(input);
            if prediction != target {
                for (w, x) in self.weights.iter_mut().zip(input.iter()) {
                    *w = *w + self.learning_rate * target_value * *x;
                }
                self.bias = self.bias + self.learning_rate * target_value;
            }
        }
    }
}

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

    #[test]
    fn test_fixed_point_arithmetic() {
        let a = Fixed::new(1 << 8, 8); // 1.0 in fixed point
        let b = Fixed::new(1 << 7, 8); // 0.5 in fixed point
        let c = a * b;
        assert_eq!(c.to_float(), 0.5);
    }

    #[test]
    fn test_matrix_operations() {
        let mut a = Matrix::new(2, 2, 8);
        let mut b = Matrix::new(2, 2, 8);
        
        a.set(0, 0, Fixed::new(1 << 8, 8));
        b.set(0, 0, Fixed::new(1 << 8, 8));
        
        let c = a.multiply(&b);
        assert_eq!(c.get(0, 0).to_float(), 1.0);
    }

    #[test]
    fn test_perceptron() {
        use models::Perceptron;
        
        let mut perceptron = Perceptron::new(2, 8);
        let input = vec![Fixed::new(1 << 8, 8), Fixed::new(1 << 8, 8)];
        
        perceptron.train(&input, true);
        assert_eq!(perceptron.predict(&input), true);
    }
}